def create_loop_op(ctx, node, batch_val_input_id, data_type): nodes = [] cond_var_name = "condition" true = helper.make_tensor(cond_var_name, TensorProto.BOOL, (), [True]) init_cond = Node( helper.make_node("Constant", [], [cond_var_name], value=true), ctx) nodes.append(init_cond) # Loop requires at least a variable, add a useless fake variable. fake_val_name = "fake_var" fake_var_init_val = helper.make_tensor(fake_val_name, TensorProto.FLOAT, (), [0.0]) fake_var_init_node = Node( helper.make_node("Constant", [], ["fake_var"], value=fake_var_init_val), ctx) nodes.append(fake_var_init_node) op_name = utils.make_name("Loop") out_name = port_name(op_name) loop_inputs = [ batch_val_input_id, # trip count cond_var_name, # termination condition fake_val_name # initial value of loop-carried dependencies ] loop_scan_output_id = port_name(op_name, 1) loop_node = Node( helper.make_node("Loop", loop_inputs, [out_name, loop_scan_output_id], name=op_name), ctx) loop_body = create_loop_body_graph(ctx, node, node.input[0], data_type) loop_node.set_attr("body", loop_body) ctx.add_body_graph(out_name, loop_body) nodes.append(loop_node) return nodes
def _create_transpose_pairs_after_node(self, node): assert len(node.output) == 1 # just support node who has 1 output non_nchw_trans_consumers = self._get_non_nchw_transpose_output_nodes( node) added_node = [] # add Transpose(0, 3, 1, 2) and Transpose(0, 2, 3, 1) before each non_nchw_trans_consumers for consumer in non_nchw_trans_consumers: nchw_op_name = utils.make_name("Transpose") nchw_out_name = utils.port_name(nchw_op_name) kwargs = {"perm": [0, 3, 1, 2]} nchw = helper.make_node("Transpose", [node.output[0]], [nchw_out_name], name=nchw_op_name, **kwargs) nhwc_op_name = utils.make_name("Transpose") nhwc_out_name = utils.port_name(nhwc_op_name) kwargs = {"perm": [0, 2, 3, 1]} nhwc = helper.make_node("Transpose", [nchw_out_name], [nhwc_out_name], name=nhwc_op_name, **kwargs) nchw_node = Node(nchw, self._g) nhwc_node = Node(nhwc, self._g) self._g.replace_input(consumer, node.output[0], nhwc_out_name) added_node.extend([nchw_node, nhwc_node]) if added_node: self._update_graph_nodes(added_node, None, True) return added_node
def create_body_graph_for_if_branch(ctx, input_id, data_shape): data_type = ctx.get_dtype(input_id) nodes = [] op_name = utils.make_name("Gather") true_out_name = port_name(op_name) true_gather_node = helper.make_node("Gather", [input_id, "i"], [true_out_name], name=op_name) nodes.append(true_gather_node) op_name = utils.make_name("Squeeze") true_squeeze_out_name = port_name(op_name) cur_true_val_scalar_node = helper.make_node("Squeeze", [true_out_name], [true_squeeze_out_name], name=op_name, axes=[0]) nodes.append(cur_true_val_scalar_node) identity_node = helper.make_node('Identity', [true_squeeze_out_name], ['y'], name=utils.make_name("Identity")) nodes.append(identity_node) # create one output y = helper.make_tensor_value_info('y', data_type, tuple(data_shape[1:])) graph_def = helper.make_graph( nodes, 'if-body-graph', [], [y], ) return graph_def
def select_op8(ctx, node, name, args): # T output = Select(bool condition, T x, T y) # V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial) if len(node.input) == 1: raise ValueError("Select with only condition is not supported.") nodes = [] data_type = ctx.get_dtype(node.input[1]) op_name = utils.make_name("Size") out_name = port_name(op_name) batch_size_node = Node( helper.make_node("Size", [node.input[0]], [out_name], name=op_name), ctx) nodes.append(batch_size_node) nodes_to_append = create_loop_op(ctx, node, batch_size_node.output[0], data_type) nodes.extend(nodes_to_append) loop_scan_output_id = nodes[-1].output[1] ctx.copy_shape(node.output[0], loop_scan_output_id) ctx.set_dtype(node.output[0], data_type) op_name = node.name out_name = port_name(op_name) output_node = Node( helper.make_node("Identity", [loop_scan_output_id], [out_name], name=op_name), ctx) nodes.append(output_node) return nodes
def sparse_softmax_cross_entropy_with_logits_op_by_gathernd( ctx, node, name, args): nodes = [] # make subgraph to implement one_hot, idea comes from onehot_op indices_name = node.input[1] indices_shape = ctx.get_shape(indices_name) if len(indices_shape) != 1: # TODO: this works for rank=1 but tensorflow supports more than this. # Same principle should work but we need to implement our own eye. raise ValueError("onehot op: only rank1 is supported") logit_name = node.input[0] logit_dtype = ctx.get_dtype(logit_name) utils.make_sure(logit_dtype, "Dtype of {} is None".format(logit_name)) indices_dtype = ctx.get_dtype(indices_name) if indices_dtype != TensorProto.INT64: indices_cast = ctx.make_node("Cast", [indices_name], attr={"to": TensorProto.INT64}) nodes.append(indices_cast) indices_name = indices_cast.output[0] indices_size = ctx.make_node("Size", [indices_name]) indices_unsqueeze = ctx.make_node("Unsqueeze", [indices_name], attr={"axes": [1]}) zero_const = ctx.make_const(utils.make_name("zero"), np.array(0, dtype=np.int64)) one_const = ctx.make_const(utils.make_name("one"), np.array(1, dtype=np.int64)) id_name = utils.make_name("sparse_softmax_id") id_output = utils.port_name(id_name) nodes.extend( make_range(ctx, zero_const.output[0], indices_size.output[0], one_const.output[0], id_output, id_name, TensorProto.INT64)) id_unsqueeze = ctx.make_node("Unsqueeze", [id_output], attr={"axes": [1]}) indices_with_id = ctx.make_node( "Concat", [id_unsqueeze.output[0], indices_unsqueeze.output[0]], attr={"axis": 1}) log_softmax = ctx.make_node(op_type="LogSoftmax", inputs=[logit_name], dtypes=[logit_dtype]) gathernd_name = utils.make_name("sparse_softmax_gathernd") gathernd_output = utils.port_name(gathernd_name) nodes.extend( make_gathernd(ctx, log_softmax.output[0], indices_with_id.output[0], gathernd_output, gathernd_name, logit_dtype)) const_name = utils.make_name("const_negative_one") const_negative_one = ctx.make_const( const_name, np.array(-1).astype(utils.ONNX_TO_NUMPY_DTYPE[logit_dtype])) mul2 = ctx.make_node( op_type="Mul", inputs=[const_negative_one.output[0], gathernd_output]) res = ctx.make_node(op_type="Squeeze", inputs=[mul2.output[0]], outputs=[node.output[0]], attr={"axes": [1]}) nodes.extend([ indices_size, indices_unsqueeze, id_unsqueeze, indices_with_id, log_softmax, mul2, res ]) return nodes
def create_loop_body_graph(ctx, node, select_condition_input_id, select_output_data_type): nodes = [] graph_inputs = [ helper.make_tensor_value_info("i", TensorProto.INT64, (1, )), # iteration_num helper.make_tensor_value_info("cond", TensorProto.BOOL, ()), # condition helper.make_tensor_value_info("fake_var", TensorProto.FLOAT, ()) # loop-carried dependency ] # get the i'th value of "Select"'s condition op_name = utils.make_name("Gather") cond_gather_out_name = port_name(op_name) cond_gather_node = helper.make_node("Gather", [select_condition_input_id, "i"], [cond_gather_out_name], name=op_name) nodes.append(cond_gather_node) op_name = utils.make_name("Squeeze") cur_cond_val_out_name = port_name(op_name) cur_cond_val_scalar_node = helper.make_node("Squeeze", [cond_gather_out_name], [cur_cond_val_out_name], name=op_name, axes=[0]) nodes.append(cur_cond_val_scalar_node) if_node, if_node_output_id = create_if_op(ctx, node, cur_cond_val_out_name) nodes.append(if_node) identity_node = helper.make_node('Identity', [if_node_output_id], ['output'], name=utils.make_name("Identity")) nodes.append(identity_node) identity_node = helper.make_node('Identity', ["cond"], ['cond_output'], name=utils.make_name("Identity")) nodes.append(identity_node) identity_node = helper.make_node('Identity', ["fake_var"], ['fake_var_output'], name=utils.make_name("Identity")) nodes.append(identity_node) output_shape = get_hidden_size_best_effort(ctx, node) graph_outputs = [ helper.make_tensor_value_info("cond_output", TensorProto.BOOL, ()), helper.make_tensor_value_info("fake_var_output", TensorProto.FLOAT, ()), helper.make_tensor_value_info("output", select_output_data_type, output_shape[1:]) ] body_graph = helper.make_graph(nodes, "loop-body-graph", graph_inputs, graph_outputs) return body_graph
def pre_optimize_action(self): # make Reshape into a const, which then can be fused into Conv's weight for mobilenet_v1_75_192 ops = self.nodes constable_reshape_ops = [ n for n in ops if (n.type == "Reshape" and n.inputs[0].is_const() and n.inputs[1].is_const()) ] for reshape_op in constable_reshape_ops: target_t = reshape_op.inputs[0].get_tensor_value(as_list=False) target_shape = reshape_op.inputs[1].get_tensor_value(as_list=False) new_data = np.reshape(target_t, tuple(target_shape)) const_name = utils.port_name(utils.make_name("Const")) # point all children nodes inputs to the new node for output_name in reshape_op.output: for child in ops: for i, name in enumerate(child.input): if name == output_name: child.input[i] = const_name ops.append(self._g.make_const(const_name, new_data)) # need call this to make input update synced to protobuf val self._g.update_proto() ops.remove(reshape_op) self._g.set_nodes(ops) self._g.topological_sort(ops)
def rewrite_random_normal(g, ops): pattern1 = \ OpTypePattern('Add', name='output', inputs=[ OpTypePattern('Mul', name='input2', inputs=[ OpTypePattern('RandomStandardNormal', name='input1', inputs=["*"]), "Const|ConstV2" ]), "Const|ConstV2" ]) pattern2 = \ OpTypePattern('Identity', name='output', inputs=[ OpTypePattern('Identity', name='input2', inputs=[ OpTypePattern('RandomStandardNormal', name='input1', inputs=["*"]) ]) ]) pattern_list = [pattern1, pattern2] for pattern in pattern_list: matcher = GraphMatcher(pattern) match_results = list(matcher.match_ops(ops)) for match in match_results: output = match.get_op('output') if output.type == 'Add': # pattern 1 mean = output.inputs[1].get_tensor_value() else: # pattern 2 mean = 0.0 input2 = match.get_op('input2') if input2.type == 'Mul': scale = input2.inputs[1].get_tensor_value() else: scale = 1.0 dtype = g.get_dtype(output.output[0]) op_name = utils.make_name("RandomNormal") out_name = utils.port_name(op_name) rn_op = match.get_op('input1') seed = float(rn_op.get_attr('seed2').i) attr = {"mean": mean, "scale": scale, "dtype": dtype, "seed": seed} if rn_op.inputs[0].type == "Shape": shape_node = rn_op.inputs[0] new_node = g.make_node("RandomNormalLike", [shape_node.input[0]], outputs=[out_name], name=op_name, attr=attr) else: shape = g.get_shape(output.output[0]) if shape is None or -1 in shape: continue attr['shape'] = shape new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name, attr=attr) g.replace_all_inputs(output.output[0], new_node.output[0], ops=ops) g.safe_remove_nodes(match.get_nodes()) return ops
def pre_optimize_action(self): # make Reshape into a const, which then can be fused into Conv's weight for mobilenet_v1_75_192 ops = self.nodes constable_reshape_ops = [ n for n in ops if (n.type == "Reshape" and self._g.is_initializer(n.input[0]) and self._g.is_initializer(n.input[1])) ] for reshape_op in constable_reshape_ops: target_t = numpy_helper.to_array( self._g.get_initializer(reshape_op.input[0])) target_shape = numpy_helper.to_array( self._g.get_initializer(reshape_op.input[1])) new_data = np.reshape(target_t, tuple(target_shape)) const_name = utils.port_name(utils.make_name("Const")) new_tensor = numpy_helper.from_array(new_data, const_name) # point all children nodes inputs to the new node for output_name in reshape_op.output: for child in ops: for i, name in enumerate(child.input): if name == output_name: child.input[i] = const_name self._g.add_initializer(new_tensor) # need call this to make input update synced to protobuf val self._g.update_proto() ops.remove(reshape_op) self._g.set_nodes(ops) self._g.topological_sort(ops)
def test_rewrite_subgraph(self): graph_proto = self.sample_net() g = GraphUtil.create_graph_from_onnx_graph(graph_proto) pattern = \ OpTypePattern('Abs', name='output', inputs=[ OpTypePattern('Add', name='input') ]) ops = g.get_nodes() matcher = GraphMatcher(pattern) match_results = list(matcher.match_ops(ops)) for match in match_results: input_node = match.get_op('input') output_node = match.get_op('output') op_name = utils.make_name("ReplacedOp") out_name = utils.port_name(op_name) new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name) g.replace_all_inputs(output_node.output[0], new_node.output[0]) # ops=ops for n in set(match.get_nodes()): g.remove_node(n.name) g.topological_sort(ops) result = onnx_to_graphviz(g) expected = 'digraph { Placeholder__4 [op_type=Placeholder] n1 [op_type=Abs] ' \ 'n3 [op_type=Abs] n2 [op_type=Abs] ReplacedOp__5 [op_type=Sub] ' \ 'n6 [op_type=Identity] n5_graph_outputs_Identity__3 [op_type=Identity] ' \ 'input -> n1 n1:0 -> n3 n1:0 -> n2 n2:0 -> ReplacedOp__5 n3:0 -> ReplacedOp__5 ' \ 'ReplacedOp__5:0 -> n6 ReplacedOp__5:0 -> n5_graph_outputs_Identity__3 }' self.assertEqual(expected, result)
def _switch_transpose_and_node(self, node, trans): if not self._transpose_has_single_consumer_node([trans]): return False input_index = 0 for i in node.input: if i == trans.output[0]: break else: input_index += 1 ops = self._g.get_nodes() self._g.replace_all_inputs(ops, node.output[0], trans.output[0]) node.input[input_index] = trans.input[0] trans.input[0] = utils.port_name(node.name) # need to transpose node shape in backward direction as well after switch # otherwise, reshape added in post_optimize_action may not work correctly shape = self._g.get_shape(node.output[0]) if shape: # only nhwc transpose can reach here new_shape = [shape[i] for i in [0, 3, 1, 2]] self._g.set_shape(node.output[0], new_shape) self._g.set_nodes(ops) return True
def test_tensor_data(self): tensors = { "empty_tensor": np.array([], dtype=np.float32), "multi_dim_empty_tensor": np.array([[], []], dtype=np.float32), "scalar": np.array(1., dtype=np.float32), "one_item_array": np.array([1.], dtype=np.float32), "normal_array": np.array([[1., 2.], [2., 3.]], dtype=np.float32) } tf_reset_default_graph() with tf_session() as sess: for n, data in tensors.items(): tf.constant(data, dtype=tf.float32, name=n) for tf_node in sess.graph.get_operations(): name = tf_node.name self.assertTrue(name in tensors.keys()) self.assertTrue("value" in tf_node.node_def.attr) # convert to onnx tensor value tensor_value = tf_utils.tf_to_onnx_tensor( tf_utils.get_tf_node_attr(tf_node, "value"), name=utils.port_name(tf_node.name) ) attr = helper.make_attribute("value", tensor_value) # same as node.get_tensor_value(is_list=False) actual = numpy_helper.to_array(helper.get_attribute_value(attr)) expected = tensors[name] self.assertTrue(np.array_equal(expected, actual))
def insert_new_node_on_output(self, op_type, output_name, name, domain=None, **kwargs): """Create and insert a new node into the graph. Args: op_type: type for new operation output_name: the names of the outputs above us name: the name of the new op kwargs: attributes of the new node Returns: node that was inserted """ utils.make_sure(isinstance(output_name, six.text_type), "output_name's type is not expected: %s", type(output_name)) utils.make_sure(isinstance(op_type, six.text_type), "op_type's type is not expected: %s", type(op_type)) new_output = port_name(name) new_node = self.make_node(op_type, [output_name], attr=kwargs, outputs=[new_output], name=name, domain=domain) to_replace = [n for n in self.get_nodes() if n != new_node] self.replace_all_inputs(to_replace, output_name, new_output) return new_node
def rewrite_dropout(g, ops): pattern = \ OpTypePattern('Mul', name='outputs', inputs=[ OpTypePattern('RealDiv', name="input2"), OpTypePattern('Floor', inputs=[ OpTypePattern('Add', inputs=[ OpTypePattern(None, name="input3"), OpTypePattern('RandomUniform|RandomUniformLike'), ]) ]), ]) matcher = GraphMatcher(pattern) match_results = list(matcher.match_ops(ops)) for match in match_results: inputs2 = match.get_op('input2') outputs = match.get_op('outputs') op_name = utils.make_name("Dropout") out_name = port_name(op_name) new_node = g.make_node("Dropout", [inputs2.input[0]], outputs=[out_name], name=op_name, attr={"ratio": 1.0}, shapes=[g.get_shape(inputs2.input[0])], dtypes=[g.get_dtype(inputs2.input[0])]) g.replace_all_inputs(ops, outputs.output[0], new_node.output[0]) g.safe_remove_nodes(match.get_nodes()) # remove dropout if its ratio is 1.0 for node in g.get_nodes(): if node.type == "Dropout" and node.get_attr("ratio").f == 1.0: g.replace_all_inputs(g.get_nodes(), node.output[0], node.input[0]) g.remove_node(node.name) return ops
def rewrite_random_normal(g, ops): pattern = \ OpTypePattern('Add', name='output', inputs=[ OpTypePattern('Mul', name='input2', inputs=[ OpTypePattern('RandomStandardNormal', name='input1', inputs=["*"]), "*" ]), "*" ]) matcher = GraphMatcher(pattern) match_results = list(matcher.match_ops(ops)) for match in match_results: output = match.get_op('output') mean = output.inputs[1].get_tensor_value() dtype = g.get_dtype(output.output[0]) op_name = utils.make_name("RandomNormal") out_name = port_name(op_name) rn_op = match.get_op('input1') if rn_op.inputs[0].type == "Shape": shape_node = rn_op.inputs[0] new_node = g.make_node("RandomNormalLike", [shape_node.input[0]], outputs=[out_name], name=op_name, attr={"mean": mean, "scale": 1.0, "dtype": dtype}) else: shape = g.get_shape(output.output[0]) new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name, attr={"shape": shape, "mean": mean, "scale": 1.0, "dtype": dtype}) g.replace_all_inputs(ops, output.output[0], new_node.output[0]) for n in set(match.get_nodes()): g.remove_node(n.name) return ops
def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=None, **kwargs): """Create and insert a new node into the graph. Args: node: we want to replace the input for this node op_type: type for new operation input_name: the names of the outputs above us name: the name of the new op kwargs: attributes of the new node Returns: node that was inserted """ if name is None: name = utils.make_name(node.name) new_output = port_name(name) new_node = self.make_node(op_type, [input_name], attr=kwargs, outputs=[new_output], name=name, domain=domain) for i, n in enumerate(node.input): if n == input_name: node.input[i] = new_output break return new_node
def replace_subgraph(ops, subgraph_nodes, old_inputs, old_outputs, new_inputs, new_outputs): """Replace subgraph.""" if len(old_inputs) != len(new_inputs) or len(old_outputs) != len( new_outputs): raise ValueError( "replace_subgraph - inputs and outputs need to be same length") # point all children nodes inputs to the new node for oo, no in zip(old_outputs, new_outputs): for output_name in oo.output: for child in ops: for i, name in enumerate(child.input): if name == output_name: child.input[i] = port_name(no.name) # delete nodes no longer used removed = set() for node in subgraph_nodes.get_nodes(): if not node or node in removed: continue ops.remove(node) removed.add(node) ops.extend(new_outputs) return ops
def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=None, extra_opset=None, output_names=None): """Create Graph. Args: nodes: list of Node() output_shapes: dict of tensorflow output shapes dtypes: dict of tensorflow dtype """ if target is None: target = [] self._nodes = [] self._nodes_by_name = {} self._output_to_node_name = {} self.shapes = {} self._target = set(target) self._dtypes = dtypes self._output_shapes = output_shapes self._opset = find_opset(opset) if extra_opset is not None: utils.make_sure(isinstance(extra_opset, list), "invalid extra_opset") self._extra_opset = extra_opset self._order_sensitive_inputs = [] self.outputs = output_names if output_names is not None else [] self.parent_graph = None self.contained_graphs = {} # {node_name: {node_attribute_name: Graph}} ops = [Node(node, self) for node in nodes] self.reset_nodes(ops) # add identity node after each output, in case it is renamed during conversion. for o in self.outputs: n = self.get_node_by_output_in_current_graph(o) new_output_name = port_name(n.name + "_" + utils.make_name("raw_output_")) n_shapes = n.output_shapes n_dtypes = n.output_dtypes body_graphs = n.graph.contained_graphs.pop(n.name, None) self.remove_node(n.name) new_outputs = [output if output != o else new_output_name for output in n.output] # domain should be passed to new node new_node = self.make_node(n.type, n.input, outputs=new_outputs, attr=n.attr, name=n.name, skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes, domain=n.domain) if body_graphs: for attr_name, body_graph in body_graphs.items(): body_graph.parent_graph = self new_node.set_body_graph_as_attr(attr_name, body_graph) self.replace_all_inputs(self.get_nodes(), o, new_output_name) self.make_node("Identity", [new_output_name], outputs=[o], op_name_scope=n.name + "_" + "graph_outputs") self.copy_shape(new_output_name, o) self.copy_dtype(new_output_name, o)
def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=None, extra_opset=None, output_names=None): """Create Graph. Args: nodes: list of Node() output_shapes: dict of tensorflow output shapes dtypes: dict of tensorflow dtype """ if target is None: target = [] self._nodes = [] self._initializers = {} self._nodes_by_name = {} self._output_to_node_name = {} self.shapes = {} self._target = set(target) self._dtypes = dtypes self._output_shapes = output_shapes self._opset = find_opset(opset) self._extra_opset = extra_opset self.inputs = [] self.outputs = output_names self.parent_graph = None self.contained_graphs = {} # {node_name: {node_attribute_name: Graph}} ops = [Node(node, self) for node in nodes] # add identity node after each output, in case it is renamed during conversion. if self.outputs: to_append = [] for n in ops: raw_outputs = n.output new_output_base_name = None index_out = 0 for i, o in enumerate(raw_outputs): if o in output_names: if not new_output_base_name: new_output_base_name = utils.make_name("raw_output_") new_out = port_name(new_output_base_name, index_out) self.replace_all_inputs(ops, o, new_out) n.output[i] = new_out index_out += 1 new_output_node = self.make_node("Identity", [new_out], outputs=[o]) to_append.append(new_output_node) self.copy_shape(o, new_out) self.set_dtype(new_out, self.get_dtype(o)) self.set_node_by_name(n) ops.extend(to_append) self.set_nodes(ops)
def create_loop_op(gather_input_ids, output_type, output_shape, trip_count_input_ids, rank): nodes = [] cond_var_name = utils.make_name("condition") true = helper.make_tensor(cond_var_name, TensorProto.BOOL, (), [True]) init_cond = helper.make_node("Constant", [], [cond_var_name], value=true, name=cond_var_name) nodes.append(init_cond) # Loop requires at least a variable, add a useless fake variable. fake_val_name = utils.make_name("fake_var") fake_var_init_val = helper.make_tensor(fake_val_name, TensorProto.FLOAT, (), [0.0]) fake_var_init_node = helper.make_node("Constant", [], [fake_val_name], value=fake_var_init_val, name=fake_val_name) nodes.append(fake_var_init_node) if rank < 1: raise ValueError("rank is < 1") trip_count_input_id = trip_count_input_ids[-1 * rank] op_name = utils.make_name("loop") fake_var_output_id = port_name(op_name) loop_inputs = [ trip_count_input_id, # trip count cond_var_name, # termination condition fake_val_name # initial value of loop-carried dependencies ] loop_scan_output_id = port_name(op_name, 1) loop_body = create_loop_body_graph(gather_input_ids, output_type, output_shape, trip_count_input_ids, rank, op_name) loop_node = helper.make_node("Loop", loop_inputs, [fake_var_output_id, loop_scan_output_id], name=op_name, body=loop_body) nodes.append(loop_node) return nodes
def get_inputs_for_current_iteration(input_id, iter_index): nodes = [] op_name = utils.make_name("Gather") cond_gather_out_name = port_name(op_name) cond_gather_node = helper.make_node("Gather", [input_id, iter_index], [cond_gather_out_name], name=op_name) nodes.append(cond_gather_node) op_name = utils.make_name("Squeeze") cur_cond_val_out_name = port_name(op_name) cur_cond_val_scalar_node = helper.make_node("Squeeze", [cond_gather_out_name], [cur_cond_val_out_name], name=op_name, axes=[0]) nodes.append(cur_cond_val_scalar_node) return nodes, cur_cond_val_out_name
def create_if_op(g, input_ids, output_data_type, output_shape): op_name = utils.make_name("If") true_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[1], op_name) false_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[2], op_name) out_name = utils.port_name(op_name) # output a scalar branches = {"then_branch": true_graph, "else_branch": false_graph} if_node = g.make_node("If", [input_ids[0]], outputs=[out_name], name=op_name, skip_conversion=True, branches=branches) return if_node, out_name
def create_if_op(g, input_ids, output_data_type, output_shape): op_name = utils.make_name("If") true_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[1], op_name) false_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[2], op_name) out_name = utils.port_name(op_name) # output a scalar if_node = g.make_node("If", [input_ids[0]], outputs=[out_name], name=op_name, skip_conversion=False) if_node.set_body_graph_as_attr("then_branch", true_graph) if_node.set_body_graph_as_attr("else_branch", false_graph) return if_node, out_name
def _add_handler(self, trans, node): if self._g.is_initializer(node.input[1]): t_p = trans.inputs[0] if t_p.type in ("Conv", "ConvTranspose") and len(t_p.input) == 2: # if Conv or ConvTranspose's bias input is not set, then we set, otherwise, we don't set # todo: maybe we can add already set bias with the input??? try later conv_inputs = [t_p.input[0], t_p.input[1], node.input[1]] conv_node = self._g.make_node(t_p.type, conv_inputs, attr=t_p.attr_onnx) ops = self._g.get_nodes() trans.input[0] = utils.port_name(conv_node.name) self._g.replace_all_inputs(ops, node.output[0], trans.output[0]) self._update_graph_nodes([conv_node], [t_p, node], True) return True return False return self._handle_node_having_branches(node)
def create_if_op(ctx, node, cur_cond_val_out_name): true_graph = create_body_graph_for_if_branch(ctx, node.input[1]) false_graph = create_body_graph_for_if_branch(ctx, node.input[2]) op_name = utils.make_name("If") out_name = port_name(op_name) # output a scalar if_node = helper.make_node("If", [cur_cond_val_out_name], [out_name], name=op_name, then_branch=true_graph, else_branch=false_graph) ctx.add_body_graph(out_name, true_graph) ctx.add_body_graph(out_name, false_graph) return if_node, out_name
def follow_inputs(self, node, num, space=""): """Follow inputs for (helpful for debugging).""" val = [] top = space == "" if num == 0: return [] val.append("{}{} {} {}".format(space, node.type, node.name, self.get_shape(port_name(node.name)))) space += " " for j in node.inputs: val.extend(self.follow_inputs(j, num - 1, space)) if top: print("\n".join(reversed(val))) print() return [] return val
def _create_transpose_pairs_before_node(self, node): non_nhwc_trans_inputs = [] for input_id, n in zip(node.input, node.inputs): if not is_nhwc_transpose(n): # check in case node has two inputs coming from a same node output. if [input_id, n] not in non_nhwc_trans_inputs: non_nhwc_trans_inputs.append([input_id, n]) added_node = [] # add Transpose(0, 3, 1, 2) and Transpose(0, 2, 3, 1) before each non_nhwc_trans_consumers for input_id, n in non_nhwc_trans_inputs: nchw_op_name = utils.make_name("Transpose") nchw_out_name = utils.port_name(nchw_op_name) kwargs = {"perm": [0, 3, 1, 2]} nchw = helper.make_node("Transpose", [input_id], [nchw_out_name], name=nchw_op_name, **kwargs) nhwc_op_name = utils.make_name("Transpose") nhwc_out_name = utils.port_name(nhwc_op_name) kwargs = {"perm": [0, 2, 3, 1]} nhwc = helper.make_node("Transpose", [nchw_out_name], [nhwc_out_name], name=nhwc_op_name, **kwargs) nchw_node = Node(nchw, self._g) nhwc_node = Node(nhwc, self._g) self._g.replace_input(node, input_id, nhwc_out_name) added_node.extend([nchw_node, nhwc_node]) if added_node: self._update_graph_nodes(added_node, None, True) return added_node
def insert_new_node_on_output(self, op_type, output_name, name=None, **kwargs): """Create and insert a new node into the graph. Args: op_type: type for new operation output_name: the names of the outputs above us name: the name of the new op kwargs: attributes of the new node Returns: node that was inserted """ assert isinstance(output_name, six.text_type) and isinstance(op_type, six.text_type) new_output = port_name(name) new_node = Node(helper.make_node(op_type, [output_name], [new_output], name=name, **kwargs), self) self.replace_all_inputs(self.get_nodes(), output_name, new_output) return new_node
def create_if_op(ctx, node, cur_cond_val_out_name): data_shape = get_hidden_size_best_effort(ctx, node) true_graph = create_body_graph_for_if_branch(ctx, node.input[1], data_shape) false_graph = create_body_graph_for_if_branch(ctx, node.input[2], data_shape) op_name = utils.make_name("If") out_name = port_name(op_name) # output a scalar if_node = helper.make_node("If", [cur_cond_val_out_name], [out_name], name=op_name, then_branch=true_graph, else_branch=false_graph) return if_node, out_name
def create_if_op(input_ids, output_data_type, output_shape): op_name = utils.make_name("If") true_graph = create_body_graph_for_if_branch(output_data_type, output_shape, input_ids[1], op_name) false_graph = create_body_graph_for_if_branch(output_data_type, output_shape, input_ids[2], op_name) out_name = port_name(op_name) # output a scalar if_node = helper.make_node("If", [input_ids[0]], [out_name], name=op_name, then_branch=true_graph, else_branch=false_graph) return if_node, out_name