Exemple #1
0
    def caffe2_net_to_onnx_graph(cls,
                                 predict_net,
                                 init_net=None,
                                 value_info=None):
        if value_info is None:
            value_info = {}
        if not isinstance(value_info, dict):
            raise ValueError('Please pass value_info as a '
                             'name -> (type, shape) dictionary')

        cls._filter_fake_init(init_net, value_info)
        cls._ssa_rewrite(predict_net, init_net, value_info)

        if init_net:
            initializer = cls.caffe2_init_net_to_initializer(init_net)
            value_info.update({
                init.name: (init.data_type, init.dims)
                for init in initializer
            })
        else:
            initializer = []

        # Check whether we have got type shape info of all input
        missing = (set(list(predict_net.external_input)) -
                   set(value_info.keys()))
        if missing:
            raise RuntimeError(
                'Could not find value info of inputs: {}'.format(
                    ', '.join(missing)))

        inputs = {}
        for name in predict_net.external_input:
            elem_type, shape = value_info[name]
            inputs[name] = np.random.randn(*shape).astype(
                mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type])

        ws, outputs = c2_native_run_net(init_net, predict_net, inputs)

        for name in predict_net.external_output:
            output = outputs[name]
            elem_type = mapping.NP_TYPE_TO_TENSOR_TYPE[output.dtype]
            shape = output.shape
            value_info[name] = (elem_type, shape)

        graph_def = GraphProto()
        graph_def.name = predict_net.name
        graph_def.initializer.extend(initializer)
        # This is a mapping from Caffe2 names to ONNX names
        graph_def.input.extend(
            make_tensor_value_info(name=name,
                                   elem_type=value_info[name][0],
                                   shape=value_info[name][1])
            for name in predict_net.external_input)

        cls._dummy_name.reset(
            cls._all_names_in_net(predict_net)
            | cls._all_names_in_net(init_net))

        for op in predict_net.op:
            shapes = {}
            for name in itertools.chain(op.input, op.output):
                blob = ws.FetchBlob(name)
                if hasattr(blob, 'shape'):
                    shapes[name] = blob.shape
            nodes, const_tensors = cls.caffe2_op_to_onnx_node(op,
                                                              shapes=shapes)
            graph_def.node.extend(nodes)
            graph_def.initializer.extend(const_tensors)
            graph_def.input.extend(
                [cls._extract_value_info(tensor) for tensor in const_tensors])

        all_output = set(
            sum((list(node.output) for node in graph_def.node),
                [init.name for init in graph_def.initializer]))
        redundant_output = set(vi.name for vi in graph_def.output) - all_output
        if redundant_output:
            logger.warning(
                'There are graph output not produced by any node or initializer: {}'
                '! Will drop them.'.format(', '.join(redundant_output)))
        graph_def.output.extend(
            make_tensor_value_info(name=name,
                                   elem_type=value_info[name][0],
                                   shape=value_info[name][1])
            for name in predict_net.external_output if name in all_output)

        return graph_def
Exemple #2
0
    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: Text,
            fname: Text,
            inputs: List[Text],
            outputs: List[Text],
            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])
Exemple #3
0
    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))
        if node_proto.output is None:
            node_proto.output = []
            output_types = output_arr.types
            type_attr = onnx.helper.make_attribute(
                output_arr.name + '-types', [
                    str(data_type).replace('tensor(', '').replace(')', '')
                    for data_type in output_types
                ])
            node_proto.attribute.append(type_attr)
        node_proto.output.append(output_arr.name)
    return node_proto


nodes = [
    create_node_from_schema(schema)
    for schema in sorted(schemas, key=lambda s: s.name)
]

with open('onnx-op-defs.pb', 'wb') as f:
    graph_proto = GraphProto()
    graph_proto.node.extend(nodes)
    f.write(graph_proto.SerializeToString())
    # for node in nodes:
    #     message_to_string = text_format.MessageToString(node, as_utf8=True)
    #     node_2 = load_node(message_to_string)
    #     f.write(message_to_string + '----f\n')

# with open('onnx.pbtxt','r') as f:
#     nodes = [load_node(node_str) for node_str in f.read().split('----f\n')]
#     print(nodes)
Exemple #5
0
    def graph_def_to_onnx_graph(
        cls,
        graph_def,
        init_func=None,
        constants=None,
        value_info=None,
        graph_name=None,
        verbose=True,
        enforce_no_running=False,
    ):
        if value_info is None: value_info = {}
        if not isinstance(value_info, dict):
            raise ValueError('Please pass value_info as a '
                             'name -> (type, shape) dictionary')

        leaf_tensors = extract_leaf_tensors(graph_def)
        initializer = extract_initializer(graph_def)

        # Check whether we have got type shape info of all input
        missing = (leaf_tensors - set(value_info.keys()) - initializer)
        if missing:
            raise RuntimeError(
                'Could not find value info of inputs: {}'.format(
                    ', '.join(missing)))

        # Check if value_info contains the types/shapes of all the blobs, in
        # which case we don't need to infer them by running the net.
        run_native_graph = False
        for op in graph_def.op:
            for name in itertools.chain(op.input, op.output):
                if name not in value_info:
                    run_native_graph = True
                    break

        ws = None

        # Get the value info of outputs and initializer
        if run_native_graph and not enforce_no_running:
            inputs = {}
            for name, (elem_type, shape) in value_info.items():
                inputs[name] = np.random.randn(*shape).astype(
                    mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type])
            ws, outputs, initializer = native_run_graph(
                graph_def, inputs, initializer, init_func)

        if enforce_no_running:
            # In some cases(e.g. PyTorch), we had ran the graph
            # outputs had been in ``value_info`` already
            import dragon.core.workspace as ws
            initializer = fetch_initializer(initializer)

        # Prepare to make the graph
        onnx_graph = GraphProto()
        onnx_graph.name = graph_name if graph_name else graph_def.name

        # Initializer should also be included in the inputs
        value_info.update(
            {init.name: (init.data_type, init.dims)
             for init in initializer})

        # Add initializer
        onnx_graph.initializer.extend(initializer)

        # Add inputs
        onnx_graph.input.extend(
            make_tensor_value_info(name=name,
                                   elem_type=value_info[name][0],
                                   shape=value_info[name][1])
            for name in leaf_tensors)

        # Add outputs
        onnx_graph.output.extend(
            make_tensor_value_info(name=name,
                                   elem_type=value_info[name][0],
                                   shape=value_info[name][1])
            for name in set(graph_def.output))

        # Add constants
        if constants is not None:
            for k, v in constants.items():
                onnx_graph.initializer.extend(
                    [numpy_helper.from_array(v, name=k)])

        # Add nodes
        shapes, ssa_names, ssa_outputs = {}, {}, defaultdict(int)

        for op in graph_def.op:
            # Get the shape of inputs and outputs
            for name in itertools.chain(op.input, op.output):
                if ws and ws.HasTensor(name):
                    blob = ws.FetchTensor(name)
                    if hasattr(blob, 'shape'):
                        shapes[name] = blob.shape
                else:
                    shapes[name] = value_info[name][1]

            # SSA rewritten
            op, shapes, ssa_names, ssa_outputs = \
                cls._ssa_rewrite(op, shapes, ssa_names, ssa_outputs)

            # Try to translate op => nodes
            nodes, const_tensors = get_nodes_def(op, shapes, ws)

            # Directly convert outputs as const tensors if necessary
            if None in nodes:
                const_tensors = [
                    numpy_helper.from_array(ws.FetchTensor(name), name=name)
                    for name in op.output
                ]
            else:
                onnx_graph.node.extend(nodes)

            # Add const tensors
            if const_tensors is not None:
                onnx_graph.initializer.extend(const_tensors)
                onnx_graph.input.extend([
                    cls._extract_value_info(tensor) for tensor in const_tensors
                ])

        if verbose: print(printable_graph(onnx_graph))

        return onnx_graph
    def convert_mx2onnx_graph(self,
                              output_idx,
                              mx_graph,
                              mx_weights,
                              in_shape,
                              in_type,
                              log=False):
        print("\nconverting weights from MxNet NDArrays to NumPy arrays.\n")
        weights = MxNetToONNXConverter.convert_weights_to_numpy(mx_weights)

        onnx_graph = GraphProto()

        initializer = []
        all_processed_nodes = []
        onnx_processed_nodes = []
        onnx_processed_inputs = []
        onnx_processed_outputs = []

        print("Output Idxes: ", output_idx)

        for idx, node in enumerate(mx_graph):
            op = node["op"]
            name = node["name"]
            if log:
                print("Converting idx: %d, op: %s, name: %s" % (idx, op, name))
            converted = MxNetToONNXConverter.convert_layer(
                node,
                mx_graph=mx_graph,
                weights=weights,
                in_shape=in_shape,
                in_type=in_type,
                proc_nodes=all_processed_nodes,
                initializer=initializer)

            if isinstance(converted, onnx_pb.ValueInfoProto):
                if idx < (len(mx_graph) - 1):
                    onnx_processed_inputs.append(converted)
                else:
                    onnx_processed_outputs.append(converted)
            elif isinstance(converted, onnx_pb.NodeProto):
                if idx not in output_idx:
                    onnx_processed_nodes.append(converted)
                else:
                    onnx_processed_nodes.append(converted)
                    onnx_processed_outputs.append(
                        make_tensor_value_info(
                            name=converted.name,
                            elem_type=mapping.NP_TYPE_TO_TENSOR_TYPE[np.dtype(
                                'float32')],
                            shape=(in_shape[0], -1)))
                    if log:
                        print("Output node is: %s" % converted.name)
            elif isinstance(converted, onnx_pb.TensorProto):
                raise ValueError("Did not expect TensorProto")
                if idx < (len(mx_graph) - 1):
                    onnx_processed_inputs.append(converted)
                else:
                    onnx_processed_outputs.append(converted)
            else:
                print(converted)
                raise ValueError("node is of an unrecognized type: %s" %
                                 type(node))

            all_processed_nodes.append(converted)

        graph = helper.make_graph(onnx_processed_nodes, "main",
                                  onnx_processed_inputs,
                                  onnx_processed_outputs)

        graph.initializer.extend(initializer)

        checker.check_graph(graph)
        return graph
Exemple #7
0
def merge_graphs(
        g1,  # type: GraphProto
        g2,  # type: GraphProto
        io_map,  # type: List[Tuple[Text, Text]]
        inputs=None,  # type: Optional[List[Text]]
        outputs=None,  # type: Optional[List[Text]]
        prefix1=None,  # type: Optional[Text]
        prefix2=None,  # type: Optional[Text]
        name=None,  # type: Optional[Text]
        doc_string=None,  # type: Optional[Text]
):  # type: (...) -> 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
Exemple #8
0
def expand_out_dim_graph(
        graph,  # type: GraphProto
        dim_idx,  # type: int
        inplace=False,  # type: Optional[bool]
):  # type: (...) -> 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
Exemple #9
0
def add_prefix_graph(
        graph,  # type: GraphProto
        prefix,  # type: Text
        rename_nodes=True,  # type: Optional[bool]
        rename_edges=True,  # type: Optional[bool]
        rename_inputs=True,  # type: Optional[bool]
        rename_outputs=True,  # type: Optional[bool]
        rename_initializers=True,  # type: Optional[bool]
        rename_value_infos=True,  # type: Optional[bool]
        inplace=False,  # type: Optional[bool]
):  # type: (...) -> 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, name):  # type: (Text, 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