from __future__ import division from __future__ import print_function from __future__ import unicode_literals from __future__ import absolute_import import collections import copy import logging import six import numpy as np from os.path import join as pathjoin from onnx import ( helper, numpy_helper, OperatorSetIdProto, AttributeProto, TensorProto, onnx_pb, ) from hetu.onnx import util from hetu.onnx.util import FindOpset from hetu.onnx import constants class Node(object): def __init__(self, node, graph=None): self._op = node self._graph = graph self._inputs = list(node.input) self._outputs = list(node.output) self._attrs = {} if graph is not None: graph.set_node_by_name(self) for a in node.attribute: self._attrs[a.name] = a @property def input_tensor_names(self): return self._inputs @input_tensor_names.setter def input_tensor_names(self, val): self._inputs = copy.deepcopy(val) @property def output_tensor_names(self): return copy.deepcopy(self._outputs) @property def name(self): return self._op.name @property def op_type(self): return self._op.op_type @op_type.setter def op_type(self, val): self._op.op_type = val @property def is_graph_input(self): return self.op_type in ['defined_in'] @property def is_graph_output(self): return self.op_type in ['defined_out'] @property def input_nodes(self): return [self._graph.get_node_by_outputname(n) for n in self._inputs] @property def op(self): return self._op def set_attr(self, name, val): self._attrs[name] = helper.make_attribute(name, val,) def get_attr(self, name, default=None): return self._attrs.get(name, default) def get_attr_value(self, name, default=None): attr = self.get_attr(name) if attr: attr_val = helper.get_attribute_value(attr) if isinstance(attr_val, bytes): attr_val = attr_val.decode('utf-8') return attr_val return default @property def is_const(self): return self.op_type in ['Const'] or \ (self.op_type in ['PlaceholderOp'] and self._attrs.get('value') is not None) def get_tensor_value(self, as_list=True): assert self.is_const, "Failed: Node {} must be Const".format(self.name) t = self.get_attr('value') t = numpy_helper.to_array(helper.get_attribute_value(t)) if as_list: t = t.tolist() return t def onnx_attrs(self): schema = util.get_schema(self.op_type, self._graph._opset) onnx_attrs = {} for name, attr in self._attrs.items(): if name == 'value': onnx_attrs[name] = self._attrs['value'] elif schema is None or schema.has_attribute(name): onnx_attrs[name] = attr return onnx_attrs def update_node_proto(self): nodes = list(self._op.input) for node in nodes: self._op.input.remove(node) self._op.input.extend(self._inputs) nodes = list(self._op.output) for node in nodes: self._op.output.remove(node) self._op.output.extend(self._outputs) del self._op.attribute[:] attr = list(self.onnx_attrs().values()) if attr: self._op.attribute.extend(attr) # add for X2hetu @property def domain(self): """Return Op type.""" return self._op.domain class Graph(object): def __init__(self, nodes, shapes=None, dtypes=None, opset=None, output_names=None): self._nodes = [] self._nodename_to_node = {} self._outputname_to_nodename = {} self._dtypes = dtypes self._shapes = shapes self._opset = FindOpset(opset) self._outputs = output_names if output_names is not None else [] ops = [Node(node, self) for node in nodes] self.update_graph_nodes(ops) def update_graph_nodes(self, ops): remained_dtypes = {} remained_shapes = {} self._outputname_to_nodename = {} for op in ops: for op_output in op.output_tensor_names: if op_output in self._dtypes: remained_dtypes[op_output] = self._dtypes[op_output] if op_output in self._shapes: remained_shapes[op_output] = self._shapes[op_output] self._outputname_to_nodename[op_output] = op.name self._nodes = ops self._nodename_to_node = {op.name: op for op in ops} self._dtypes = remained_dtypes self._shapes = remained_shapes def set_node_by_name(self, node): self._nodename_to_node[node.name] = node for outputname in node._outputs: self._outputname_to_nodename[outputname] = node.name def get_node_by_outputname(self, outputname): nodename = self._outputname_to_nodename.get(outputname) if nodename: return self._nodename_to_node.get(nodename) return None def get_shape(self, name): return self._shapes.get(name) def set_shape(self, name, val): self._shapes[name] = val def get_dtype(self, name): return self._dtypes.get(name) def set_dtype(self, name, val): self._dtypes[name] = val def update_node_shape_dtype(self, node): if node.is_const or node.is_graph_input: return initializers = [] for i, inp in enumerate(node.input_nodes): if inp.is_const: tensor = util.TensorProtoFromNumpy(inp.get_tensor_value(as_list=False), name=inp.output_tensor_names[0]) initializers.append(tensor) input_shapes = [self.get_shape(i) for i in node.input_tensor_names] input_dtypes = [self.get_dtype(i) for i in node.input_tensor_names] shapes, dtypes = util.InferOnnxShapeDtype( node, self._opset, input_shapes, input_dtypes, initializers) if not shapes or not dtypes: return for output, shape, dtype in zip(node.output_tensor_names, shapes, dtypes): self.set_dtype(output, dtype) self.set_shape(output, shape) def make_const(self, name, np_val, raw=False, is_0D_tensor=False): shape = [] if is_0D_tensor else np_val.shape if raw: onnx_tensor = None # fixme: Not yet implemented pass else: onnx_tensor = helper.make_tensor( name, util.numpy_to_onnx_dtype(np_val.dtype), shape, np_val, raw=False, ) dtype = onnx_tensor.data_type node = self.make_node( "Const", [], outputs=[name], name=name, attr={"value": onnx_tensor}, dtypes=[dtype], ) self.set_shape(name, shape) self.set_dtype(name, dtype) return node def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, name=None, shapes=None, dtypes=None): if attr is None: attr = {} if shapes is None: shapes = [] if dtypes is None: dtypes = [] if name is None: name = util.make_name(op_type) if outputs is None: outputs = [name+':'+str(i) for i in range(output_count)] output_count = len(outputs) onnx_node = helper.make_node( op_type, inputs, outputs, name=name, **attr) node = Node(onnx_node, self) if shapes: assert len( shapes) == output_count, "Failed: output shapes count not equal to output count when make_node" for i in range(output_count): self.set_shape(node._outputs[i], shapes[i]) if dtypes: assert len( dtypes) == output_count, "Failed: output dtypes count not equal to output count when make_node" for i in range(output_count): self.set_dtype(node._outputs[i], dtypes[i]) if not shapes or not dtypes: self.update_node_shape_dtype(node) self._nodes.append(node) return node def insert_new_node_on_input(self, node, op_type, input_name, name=None, **kwargs): if name is None: name = util.make_name(node.name) new_output = util.make_name(name) if not isinstance(input_name, list): input_name = [input_name] new_node = self.make_node( op_type, input_name, attr=kwargs, outputs=[new_output], name=name, ) for i, n in enumerate(node.input_tensor_names): if n == input_name[0]: node.input_tensor_names[i] = new_output break return new_node def insert_new_node_on_output(self, op_type, output_name, name, **kwargs): new_output = util.make_name(name) new_node = self.make_node( op_type, [output_name], attr=kwargs, outputs=[new_output], name=name, ) for node in self._nodes: if node == new_node: continue for i, input_name in enumerate(node.input_tensor_names): if input_name == output_name: node.input_tensor_names[i] = new_output return new_node def replace_input(self, node, old_input, new_input, input_index=None): if input_index is None: for i, input_name in enumerate(node._inputs): if input_name == old_input: node._inputs[i] = new_input elif node._inputs[input_index] == old_input: node._inputs[input_index] = new_input else: raise RuntimeError("Failed:Unable to replace input %r into %r for node %r." % ( old_input, new_input, node.name)) def topology_sort(self, ops): def _push_stack(stack, node, in_stack): stack.append(node) if node in in_stack: raise ValueError("Graph has cycles.") in_stack[node] = True def _get_unvisited_child(g, node, not_visited): for child in g[node]: if child in not_visited: return child return -1 ops.sort(key=lambda op: op.name) n = len(ops) g = [[] for _ in range(n)] op_name_to_index = {} for i, op in enumerate(ops): op_name_to_index[op.name] = i for i, op in enumerate(ops): all_input = list(op.input_tensor_names) for inp in sorted(all_input): j = self.get_node_by_outputname(inp) g[op_name_to_index[j.name]].append(i) label = [-1 for _ in range(n)] stack = [] in_stack = dict() not_visited = dict.fromkeys(range(n)) label_counter = n-1 while not_visited: node = list(not_visited.keys())[0] _push_stack(stack, node, in_stack) while stack: node = _get_unvisited_child(g, stack[-1], not_visited) if node != -1: _push_stack(stack, node, in_stack) else: node = stack.pop() in_stack.pop(node) not_visited.pop(node) label[node] = label_counter label_counter -= 1 ret = [x for _, x in sorted(zip(label, ops))] self.update_graph_nodes(ret) def make_model(self, graph_doc, onnx_filename, graph_name='hetu.python.onnx'): graph = self.make_graph( graph_doc, onnx_filename, graph_name=graph_name, ) model_proto = helper.make_model(graph) return model_proto def make_graph(self, doc, onnx_filename, graph_name='hetu.python.onnx'): for node in self._nodes: node.update_node_proto() ops = [] const_ops = [] input_ops = [] for op in self._nodes: if op.is_const: const_ops.append(op) continue if op.is_graph_input: input_ops.append(op) continue ops.append(op) initializers = [] for op in const_ops: tensor_name = op.output_tensor_names[0] tensor = util.TensorProtoFromNumpy( op.get_tensor_value(as_list=False), tensor_name, export_path=onnx_filename, ) initializers.append(tensor) # sorted inputs by input id. input_tensor_name like this: A:0,B:1 # fixme:mybe outputs should be sort also input_ids = [op.output_tensor_names[0] for op in input_ops] input_ids = sorted(input_ids, key=lambda x: int(x.split('-')[-1])) if self._opset < 9: input_ids += [op.output_tensor_names[0] for op in const_ops] input_tensor_values = self.MakeOnnxGraphIO(input_ids) output_tensor_values = self.MakeOnnxGraphIO(self._outputs) graph = helper.make_graph( [op.op for op in ops], graph_name, input_tensor_values, output_tensor_values, initializer=initializers, doc_string=doc, ) return graph def MakeOnnxGraphIO(self, ids): tensor_value_infos = [] for name in ids: dtype = self.get_dtype(name) shape = self.get_shape(name) v = util.MakeOnnxInputsOutputs(name, dtype, shape) tensor_value_infos.append(v) return tensor_value_infos def copy_shape(self, input_name, output_name): shape = self.get_shape(input_name) if shape is not None: self.set_shape(output_name, shape)