def make_graph(
nodes: Sequence[NodeProto],
name: Text,
inputs: Sequence[ValueInfoProto],
outputs: Sequence[ValueInfoProto],
initializer: Optional[Sequence[TensorProto]] = None,
doc_string: Optional[Text] = None,
value_info: Sequence[ValueInfoProto] = [],
sparse_initializer: Optional[Sequence[SparseTensorProto]] = None,
) -> GraphProto:
if initializer is None:
initializer = []
if sparse_initializer is None:
sparse_initializer = []
if value_info is None:
value_info = []
graph = GraphProto()
graph.node.extend(nodes)
graph.name = name
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
graph.sparse_initializer.extend(sparse_initializer)
graph.value_info.extend(value_info)
if doc_string:
graph.doc_string = doc_string
return graph
def make_graph(
nodes, # type: Sequence[NodeProto]
name, # type: Text
inputs, # type: Sequence[ValueInfoProto]
outputs, # type: Sequence[ValueInfoProto]
initializer=None, # type: Optional[Sequence[TensorProto]]
doc_string=None, # type: Optional[Text]
value_info=[], # type: Sequence[ValueInfoProto]
sparse_initializer=None, # type: Optional[Sequence[SparseTensorProto]]
): # type: (...) -> GraphProto
if initializer is None:
initializer = []
if sparse_initializer is None:
sparse_initializer = []
if value_info is None:
value_info = []
graph = GraphProto()
graph.node.extend(nodes)
graph.name = name
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
graph.sparse_initializer.extend(sparse_initializer)
graph.value_info.extend(value_info)
if doc_string:
graph.doc_string = doc_string
return graph
def test_attr_repeated_graph_proto(self): # type: () -> None
graphs = [GraphProto(), GraphProto()]
graphs[0].name = "a"
graphs[1].name = "b"
attr = helper.make_attribute("graphs", graphs)
self.assertEqual(attr.name, "graphs")
self.assertEqual(list(attr.graphs), graphs)
checker.check_attribute(attr)
def make_graph(nodes, name, inputs, outputs, initializer=None, doc_string=None):
if initializer is None:
initializer = []
graph = GraphProto()
graph.node.extend(nodes)
graph.name = name
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
if doc_string:
graph.doc_string = doc_string
return graph
def insert_node(graph: GraphProto, index: int, node: NodeProto):
"""
Insert helper function to insert a node at a given index in the link list of the graph.
Args:
graph (GraphProto): Onnx graph to be modify
index (int): Index to insert the new node
node (NodeProto): The node to be inserted
"""
graph.extend([graph[-1]])
for i in reversed(range(index + 1, len(graph) - 1)):
graph[i].CopyFrom(graph[i - 1])
graph[index].CopyFrom(node)
def make_graph(
nodes: Sequence[NodeProto],
name: Text,
inputs: Sequence[ValueInfoProto],
outputs: Sequence[ValueInfoProto],
initializer: Optional[Sequence[TensorProto]] = None,
doc_string: Optional[Text] = None,
value_info: Sequence[ValueInfoProto] = [],
sparse_initializer: Optional[Sequence[SparseTensorProto]] = None,
) -> GraphProto:
"""Construct a GraphProto
Arguments:
nodes: list of NodeProto
name (string): graph name
inputs: list of ValueInfoProto
outputs: list of ValueInfoProto
initializer: list of TensorProto
doc_string (string): graph documentation
value_info: list of ValueInfoProto
sparse_initializer: list of SparseTensorProto
Returns:
GraphProto
"""
if initializer is None:
initializer = []
if sparse_initializer is None:
sparse_initializer = []
if value_info is None:
value_info = []
graph = GraphProto()
graph.node.extend(nodes)
graph.name = name
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
graph.sparse_initializer.extend(sparse_initializer)
graph.value_info.extend(value_info)
if doc_string:
graph.doc_string = doc_string
return graph
def singa_to_onnx_graph(cls, inputs, y, model_name="sonnx"):
"""
get onnx model from singa computational graph
Args:
inputs: a list of input tensors (each is initialized with a name)
Args:
y: a list of tensors, usually the outputs of the graph
Returns:
the onnx model
"""
assert len(y) == 1 # assume there is only one output
y = y[0]
graph_def = GraphProto()
graph_def.name = model_name
topol = postorderRecursive(y.creator, y)
# since tensor's name might change
# we record its id
input_tensors = {id(x): x for x in inputs}
# print(input_tensors)
X = []
Y = [helper.make_tensor_value_info(y.name, TensorProto.FLOAT, y.shape)]
for op, yid, op_t in topol:
optype = cls._get_singa_op_type(op)
# print(op.name, cls._get_singa_op_type(op), op_t, optype, yid)
if yid in input_tensors and optype == 'Dummy':
# find the input by its id
op_t = input_tensors[yid]
dtype = TensorProto.FLOAT
if op_t.dtype == tensor.int32:
dtype = TensorProto.INT32
X.append(
helper.make_tensor_value_info(op.name, dtype, op_t.shape))
else:
graph_def.node.extend(cls.singa_op_to_onnx_node(op, op_t))
graph_def.input.extend(X)
graph_def.output.extend(Y)
return graph_def
def test_get_inputs(self):
model = OnnxModel(model_proto=ModelProto(
graph=GraphProto(initializer=[TensorProto(name='y')],
input=[
ValueInfoProto(name='x'),
ValueInfoProto(name='y'),
ValueInfoProto(name='z')
])),
input_data_formats=[None, None])
self.assertEqual(model.get_inputs(),
[ValueInfoProto(name='x'),
ValueInfoProto(name='z')])
def make_graph(
nodes, # type: Sequence[NodeProto]
name, # type: Text
inputs, # type: Sequence[ValueInfoProto]
outputs, # type: Sequence[ValueInfoProto]
initializer=None, # type: Optional[Sequence[TensorProto]]
doc_string=None, # type: Optional[Text]
value_info=[], # type: Sequence[ValueInfoProto]
): # type: (...) -> GraphProto
if initializer is None:
initializer = []
if value_info is None:
value_info = []
graph = GraphProto()
graph.node.extend(nodes)
graph.name = name
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
graph.value_info.extend(value_info)
if doc_string:
graph.doc_string = doc_string
return graph
def make_model(nodes: Sequence[INodeProto],
inputs: Sequence[IValueInfoProto],
outputs: Sequence[IValueInfoProto],
initializer=None) -> IModelProto:
if initializer is None:
initializer = []
graph = GraphProto()
graph.node.extend(nodes)
graph.input.extend(inputs)
graph.output.extend(outputs)
graph.initializer.extend(initializer)
model = ModelProto()
model.graph.CopyFrom(graph)
return model
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)
dummy_name(
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)
checker.check_graph(graph_def)
return graph_def
def convert_mx2onnx_graph(self,
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 = []
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_pb2.ValueInfoProto):
if idx < (len(mx_graph) - 1):
onnx_processed_inputs.append(converted)
else:
onnx_processed_outputs.append(converted)
elif isinstance(converted, onnx_pb2.NodeProto):
if idx < (len(mx_graph) - 1):
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_pb2.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
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 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
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 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
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)
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 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] = numpy.random.randn(*shape).astype(
mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type])
ws, outputs, initializer = native_run_graph(
graph_def, inputs, initializer, init_func)
for name in graph_def.output:
output = outputs[name]
elem_type = mapping.NP_TYPE_TO_TENSOR_TYPE[output.dtype]
shape = output.shape
value_info[name] = (elem_type, shape)
if enforce_no_running:
# In some cases(e.g. PyTorch), we had ran the graph
# outputs had been in ``value_info`` already
ws = _workspace.get_default_workspace()
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 _write_onnx(onnx_graph: onnx.GraphProto, out: Path) -> None:
with open(out, "wb") as fp:
fp.write(onnx_graph.SerializeToString())