def test_merge_drop_unnecessary_initializers_and_value_info(self) -> None: ''' Tests automatic removal of initializers when merging graphs ''' ops = [helper.make_opsetid("", 10)] g = GraphProto() g.input.extend( [helper.make_tensor_value_info('x', TensorProto.FLOAT, [])]) g.output.extend( [helper.make_tensor_value_info('y', TensorProto.FLOAT, [])]) g.node.extend( [helper.make_node('Identity', inputs=['x'], outputs=['y'])]) g1 = GraphProto() g1.CopyFrom(g) g1.name = 'g1' m1 = helper.make_model(g1, producer_name='test', opset_imports=ops) checker.check_model(m1) g2 = GraphProto() g2.CopyFrom(g) g2.name = 'g2' g2.initializer.extend([ helper.make_tensor(name='x', data_type=TensorProto.FLOAT, dims=(), vals=[0]) ]) m2 = helper.make_model(g2, producer_name='test', opset_imports=ops) checker.check_model(m2) g3 = GraphProto() g3.CopyFrom(g) g3.name = 'g3' g3.sparse_initializer.extend([_make_sparse_tensor('x')]) m3 = helper.make_model(g3, producer_name='test', opset_imports=ops) checker.check_model(m3) g4 = GraphProto() g4.CopyFrom(g) g4.name = 'g3' g4.value_info.extend( [helper.make_tensor_value_info('x', TensorProto.FLOAT, [])]) m4 = helper.make_model(g4, producer_name='test', opset_imports=ops) checker.check_model(m4) # Initializer 'x' from m1 is removed, because there is no longer an input with that name out_m1 = compose.merge_models(m1, m2, prefix1='m1/', io_map=[('y', 'x')]) self.assertEqual(0, len(out_m1.graph.initializer)) # Sparse initializer 'x' from m1 is removed, because there is no longer an input with that name out_m2 = compose.merge_models(m1, m3, prefix1='m1/', io_map=[('y', 'x')]) self.assertEqual(0, len(out_m2.graph.initializer)) # Value info 'x' from m1 is removed, because there is no longer an input with that name out_m3 = compose.merge_models(m1, m4, prefix1='m1/', io_map=[('y', 'x')]) self.assertEqual(0, len(out_m3.graph.value_info))
def test_overlapping_function_names(self) -> None: ''' Tests error checking when the name of local function entries overlaps ''' ops = [helper.make_opsetid("", 10), helper.make_opsetid("local", 10)] def _make_function( domain: str, fname: str, inputs: List[str], outputs: List[str], nodes: List[NodeProto], ) -> FunctionProto: f = FunctionProto() f.domain = domain f.name = fname f.input.extend(inputs) f.output.extend(outputs) f.node.extend(nodes) f.opset_import.extend(ops) return f ops = [helper.make_opsetid("", 10), helper.make_opsetid("local", 10)] g = GraphProto() g.input.extend([ helper.make_tensor_value_info('x0', TensorProto.FLOAT, []), helper.make_tensor_value_info('x1', TensorProto.FLOAT, []) ]) g.output.extend([ helper.make_tensor_value_info('y', TensorProto.FLOAT, []), ]) g.node.extend([ helper.make_node('f1', domain='local', inputs=['x0', 'x1'], outputs=['y']) ]) g1 = GraphProto() g1.CopyFrom(g) g1.name = 'g1' m1 = helper.make_model(g1, producer_name='test', opset_imports=ops) m1.functions.extend([ _make_function( 'local', 'f1', ['x0', 'x1'], ['y'], [helper.make_node('Add', inputs=['x0', 'x1'], outputs=['y'])]) ]) checker.check_model(m1) g2 = GraphProto() g2.CopyFrom(g) g2.name = 'g2' m2 = helper.make_model(g2, producer_name='test', opset_imports=ops) m2.functions.extend([ _make_function( 'local', 'f1', ['x0', 'x1'], ['y'], [helper.make_node('Mul', inputs=['x0', 'x1'], outputs=['y'])]) ]) checker.check_model(m2) m = compose.merge_models(m1, m2, io_map=[('y', 'x0'), ('y', 'x1')], prefix1='m1/', prefix2='m2/') checker.check_model(m) nodes = [n.op_type for n in m.graph.node] self.assertEqual(['m1/f1', 'm2/f1'], nodes) functions = [f.name for f in m.functions] self.assertEqual(['m1/f1', 'm2/f1'], functions) g3 = GraphProto() g3.CopyFrom(g) g3.name = 'g3' g3.node[0].op_type = 'f2' m3 = helper.make_model(g3, producer_name='test', opset_imports=ops) m3.functions.extend([ _make_function('local', 'f1', ['x0', 'x1'], ['y'], [ helper.make_node('Add', inputs=['x0', 'x1'], outputs=['y0']), helper.make_node('Mul', inputs=['x0', 'x1'], outputs=['y1']), helper.make_node('Add', inputs=['y0', 'y1'], outputs=['y']) ]), _make_function('local', 'f2', ['x0', 'x1'], ['y'], [ helper.make_node( 'f1', domain='local', inputs=['x0', 'x1'], outputs=['y0']), helper.make_node('Mul', inputs=['x0', 'x1'], outputs=['y1']), helper.make_node('Add', inputs=['y0', 'y1'], outputs=['y']) ]) ]) checker.check_model(m3) m = compose.merge_models(m1, m3, io_map=[('y', 'x0'), ('y', 'x1')], prefix1='m1/', prefix2='m3/') checker.check_model(m) nodes = [n.op_type for n in m.graph.node] self.assertEqual(['m1/f1', 'm3/f2'], nodes) functions = [f.name for f in m.functions] self.assertEqual(['m1/f1', 'm3/f1', 'm3/f2'], functions) self.assertEqual(['Add'], [n.op_type for n in m.functions[0].node]) self.assertEqual(['Add', 'Mul', 'Add'], [n.op_type for n in m.functions[1].node]) self.assertEqual(['m3/f1', 'Mul', 'Add'], [n.op_type for n in m.functions[2].node])
def expand_out_dim_graph( graph: GraphProto, dim_idx: int, inplace: Optional[bool] = False, ) -> GraphProto: """Inserts an extra dimension with extent 1 to each output in the graph. Inserts an Unsqueeze node for each output. It can be used as a utility before merging graphs, for example when the second one expects a batch dimension. Arguments: graph (GraphProto): Graph dim_idx (int): Index of the dimension to be inserted. A negative value means counting dimensions from the back. inplace (bool): If True, mutates the model directly. Otherwise, a copy will be created """ if type(graph) is not GraphProto: raise ValueError("graph argument is not an ONNX graph") if not inplace: g = GraphProto() g.CopyFrom(graph) else: g = graph orig_out_names = [output.name for output in g.output] for n in g.node: for i in range(len(n.output)): if n.output[i] in orig_out_names: n.output[i] = n.output[i] + f'_collapsed_dim_{dim_idx}' for i in range(len(n.input)): if n.input[i] in orig_out_names: n.input[i] = n.input[i] + f'_collapsed_dim_{dim_idx}' expand_dim_k = g.name + "_expand_out_dim_idx" g.node.append( helper.make_node('Constant', inputs=[], outputs=[expand_dim_k], name=f"{expand_dim_k}-constant", value=helper.make_tensor(name=f"{expand_dim_k}-value", data_type=tp.INT64, dims=[ 1, ], vals=[ dim_idx, ]))) for _ in range(len(g.output)): o = g.output.pop(0) prev_output = o.name + f'_collapsed_dim_{dim_idx}' g.node.append( helper.make_node('Unsqueeze', inputs=[prev_output, expand_dim_k], outputs=[o.name], name=f"unsqueeze-{o.name}")) new_shape = [d.dim_value for d in o.type.tensor_type.shape.dim] new_shape.insert(dim_idx, 1) g.output.append( helper.make_tensor_value_info(o.name, o.type.tensor_type.elem_type, new_shape)) return g
def merge_graphs( g1: GraphProto, g2: GraphProto, io_map: List[Tuple[Text, Text]], inputs: Optional[List[Text]] = None, outputs: Optional[List[Text]] = None, prefix1: Optional[Text] = None, prefix2: Optional[Text] = None, name: Optional[Text] = None, doc_string: Optional[Text] = None, ) -> GraphProto: """Combines two ONNX graphs into a single one. The combined graph is defined by connecting the specified set of outputs/inputs. Those inputs/outputs not specified in the io_map argument will remain as inputs/outputs of the combined graph. Arguments: g1 (GraphProto): First graph g2 (GraphProto): Second graph io_map (list of pairs of string): The pairs of names [(out0, in0), (out1, in1), ...] representing outputs of the first graph and inputs of the second to be connected inputs (list of string): Optional list of inputs to be included in the combined graph By default, all inputs not present in the ``io_map`` argument will be included in the combined model outputs (list of string): Optional list of outputs to be included in the combined graph By default, all outputs not present in the ``io_map`` argument will be included in the combined model prefix1 (string): Optional prefix to be added to all names in g1 prefix2 (string): Optional prefix to be added to all names in g2 name (string): Optional name for the combined graph By default, the name is g1.name and g2.name concatenated with an undescore delimiter doc_string (string): Optional docstring for the combined graph If not provided, a default docstring with the concatenation of g1 and g2 docstrings is used """ if type(g1) is not GraphProto: raise ValueError("g1 argument is not an ONNX graph") if type(g2) is not GraphProto: raise ValueError("g2 argument is not an ONNX graph") # Prefixing names in the graph if requested, adjusting io_map accordingly if prefix1 or prefix2: if prefix1: g1_copy = GraphProto() g1_copy.CopyFrom(g1) g1 = g1_copy g1 = add_prefix_graph(g1, prefix=prefix1) if prefix2: g2_copy = GraphProto() g2_copy.CopyFrom(g2) g2 = g2_copy g2 = add_prefix_graph(g2, prefix=prefix2) io_map = [(prefix1 + io[0] if prefix1 else io[0], prefix2 + io[1] if prefix2 else io[1]) for io in io_map] io_map_g1_outs = set([io[0] for io in io_map]) io_map_g2_ins = set([io[1] for io in io_map]) reversed_io_map = {in_name: out_name for out_name, in_name in io_map} g1_outs = set([o.name for o in g1.output]) g2_ins = set([i.name for i in g2.input]) # If necessary extract subgraphs if inputs or outputs: if not inputs: g1_inputs = [i.name for i in g1.input] g2_inputs = [i.name for i in g2.input] else: input_set = set(inputs) g1_inputs = [i.name for i in g1.input if i.name in input_set] g2_inputs = [ i.name for i in g2.input if i.name in input_set or i.name in io_map_g2_ins ] if not outputs: g1_outputs = [o.name for o in g1.input] g2_outputs = [o.name for o in g2.input] else: output_set = set(outputs) g1_outputs = [ o.name for o in g1.output if o.name in output_set or o.name in io_map_g1_outs ] g2_outputs = [o.name for o in g2.output if o.name in output_set] if len(g1_inputs) < len(g1.input) or len(g1_outputs) < len(g1.output): e1 = utils.Extractor(helper.make_model(g1)) g1 = e1.extract_model(g1_inputs, g1_outputs).graph if len(g2_inputs) < len(g2.input) or len(g2_outputs) < len(g2.output): e2 = utils.Extractor(helper.make_model(g2)) g2 = e2.extract_model(g2_inputs, g2_outputs).graph # Check that input/output names specified in the io_map argument are valid input/output names for g1_out_name, g2_in_name in io_map: if g1_out_name not in g1_outs: raise ValueError(f"Output {g1_out_name} is not present in g1") if g2_in_name not in g2_ins: raise ValueError(f"Input {g2_in_name} is not present in g2") # Check for name collision overlapping_names = check_overlapping_names(g1, g2, io_map) if len(overlapping_names) > 0: category, names = overlapping_names[0] raise ValueError( "Cant merge two graphs with overlapping names. " f"Found repeated {category} names: " + ", ".join(names) + "\n" + "Consider using ``onnx.compose.add_prefix`` to add a prefix to names in one of the graphs." ) g = GraphProto() g.node.extend(g1.node) g2_nodes_begin = len(g.node) g.node.extend(g2.node) g2_nodes_end = len(g.node) # Connecting outputs of the first graph with the inputs of the second for node_idx in range(g2_nodes_begin, g2_nodes_end): node = g.node[node_idx] for index, name in enumerate(node.input): if name in reversed_io_map: node.input[index] = reversed_io_map[name] if inputs: input_set = set(inputs) g.input.extend([i for i in g1.input if i.name in input_set]) g.input.extend([i for i in g2.input if i.name in input_set]) else: g.input.extend(g1.input) g.input.extend([i for i in g2.input if i.name not in io_map_g2_ins]) if outputs: output_set = set(outputs) g.output.extend([o for o in g1.output if o.name in output_set]) g.output.extend([o for o in g2.output if o.name in output_set]) else: g.output.extend([o for o in g1.output if o.name not in io_map_g1_outs]) g.output.extend(g2.output) g.initializer.extend(g1.initializer) g.initializer.extend( [init for init in g2.initializer if init.name not in io_map_g2_ins]) g.sparse_initializer.extend(g1.sparse_initializer) g.sparse_initializer.extend([ init for init in g2.sparse_initializer if init.values.name not in io_map_g2_ins ]) g.value_info.extend(g1.value_info) g.value_info.extend( [vi for vi in g2.value_info if vi.name not in io_map_g2_ins]) g.name = name if name is not None else "_".join([g1.name, g2.name]) if doc_string is None: doc_string = f"Graph combining {g1.name} and {g2.name}\n" + \ g1.name + "\n\n" + g1.doc_string + "\n\n" + g2.name + "\n\n" + g2.doc_string g.doc_string = doc_string return g
def add_prefix_graph( graph: GraphProto, prefix: Text, rename_nodes: Optional[bool] = True, rename_edges: Optional[bool] = True, rename_inputs: Optional[bool] = True, rename_outputs: Optional[bool] = True, rename_initializers: Optional[bool] = True, rename_value_infos: Optional[bool] = True, inplace: Optional[bool] = False, ) -> GraphProto: """Adds a prefix to names of elements in a graph: nodes, edges, inputs, outputs, initializers, sparse initializer, value infos. It can be used as a utility before merging graphs that have overlapping names. Empty names are not prefixed. Arguments: graph (GraphProto): Graph prefix (Text): Prefix to be added to each name in the graph rename_nodes (bool): Whether to prefix node names rename_edges (bool): Whether to prefix node edge names rename_inputs (bool): Whether to prefix input names rename_outputs (bool): Whether to prefix output names rename_initializers (bool): Whether to prefix initializer and sparse initializer names rename_value_infos (bool): Whether to prefix value info names inplace (bool): If True, mutates the graph directly. Otherwise, a copy will be created """ if type(graph) is not GraphProto: raise ValueError("graph argument is not an ONNX graph") if not inplace: g = GraphProto() g.CopyFrom(graph) else: g = graph def _prefixed(prefix: Text, name: Text) -> Text: return prefix + name if len(name) > 0 else name name_map = {} if rename_edges: for n in g.node: for e in n.input: name_map[e] = _prefixed(prefix, e) for e in n.output: name_map[e] = _prefixed(prefix, e) else: if rename_outputs: for entry in g.output: name_map[entry.name] = _prefixed(prefix, entry.name) if rename_inputs: for entry in g.input: name_map[entry.name] = _prefixed(prefix, entry.name) if rename_nodes: for n in g.node: n.name = _prefixed(prefix, n.name) if rename_initializers: for init in g.initializer: name_map[init.name] = _prefixed(prefix, init.name) for sparse_init in g.sparse_initializer: name_map[sparse_init.values.name] = _prefixed( prefix, sparse_init.values.name) name_map[sparse_init.indices.name] = _prefixed( prefix, sparse_init.indices.name) if rename_value_infos: for entry in g.value_info: name_map[entry.name] = _prefixed(prefix, entry.name) for n in g.node: for i in range(len(n.output)): if n.output[i] in name_map: n.output[i] = name_map[n.output[i]] for i in range(len(n.input)): if n.input[i] in name_map: n.input[i] = name_map[n.input[i]] for in_desc in g.input: if in_desc.name in name_map: in_desc.name = name_map[in_desc.name] for out_desc in g.output: if out_desc.name in name_map: out_desc.name = name_map[out_desc.name] for initializer in g.initializer: if initializer.name in name_map: initializer.name = name_map[initializer.name] for sparse_initializer in g.sparse_initializer: if sparse_initializer.values.name in name_map: sparse_initializer.values.name = name_map[ sparse_initializer.values.name] if sparse_initializer.indices.name in name_map: sparse_initializer.indices.name = name_map[ sparse_initializer.indices.name] for value_info in g.value_info: if value_info.name in name_map: value_info.name = name_map[value_info.name] return g