def expand_out_dim(
model: ModelProto,
dim_idx: int,
inplace: Optional[bool] = False,
) -> ModelProto:
"""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:
model (ModelProto): Model
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(model) is not ModelProto:
raise ValueError("model argument is not an ONNX model")
if not inplace:
m = ModelProto()
m.CopyFrom(model)
model = m
expand_out_dim_graph(
model.graph,
dim_idx,
inplace=True # No need to create a copy, since it's a new model
)
return model
def test_expand_out_dim(self) -> None:
'''
Tests expanding output dimensions. The resulting graph should have the same output names,
but with one more dimension at the specified index.
'''
m1 = _load_model(m1_def)
def _check_model(m1: ModelProto, m2: ModelProto, dim_idx: int) -> None:
for out_g2, out_g1 in zip(m2.graph.output, m1.graph.output):
self.assertEqual(out_g2.name, out_g1.name)
self.assertEqual(out_g2.type.tensor_type.elem_type,
out_g1.type.tensor_type.elem_type)
expected_out_shape = _get_shape(out_g1)
expected_out_shape.insert(dim_idx, 1)
self.assertEqual(_get_shape(out_g2), expected_out_shape)
for dim_idx in [0, 2, -1, -3]:
m2 = compose.expand_out_dim(m1, dim_idx)
_check_model(m1, m2, dim_idx)
# Test inplace
m2 = ModelProto()
m2.CopyFrom(m1)
dim_idx = 0
compose.expand_out_dim(m2, dim_idx, inplace=True)
_check_model(m1, m2, dim_idx)
def __init__(self, model: onnx.ModelProto):
copy_model = ModelProto()
copy_model.CopyFrom(model)
self.model = copy_model
self.input_tensors = get_input_tensors(model)
self.value_info = {
vi.name: vi
for vi in list(self.model.graph.value_info) +
list(self.model.graph.input) + list(self.model.graph.output)
}
# raise Exception if input_tensor is not defined in model.graph.input
for input_tensor in self.input_tensors:
if input_tensor not in [
input.name for input in self.model.graph.input
]:
raise Exception(
'input_tensor: %s is not defined in model.graph.input' %
input_tensor)
def _test_add_prefix(self,
rename_nodes: bool = False,
rename_edges: bool = False,
rename_inputs: bool = False,
rename_outputs: bool = False,
rename_initializers: bool = False,
rename_value_infos: bool = False,
inplace: bool = False) -> None:
m1 = _load_model(m1_def)
prefix = 'pre/'
if inplace:
m2 = ModelProto()
m2.CopyFrom(m1)
compose.add_prefix(m2,
prefix,
rename_nodes=rename_nodes,
rename_edges=rename_edges,
rename_inputs=rename_inputs,
rename_outputs=rename_outputs,
rename_initializers=rename_initializers,
rename_value_infos=rename_value_infos,
inplace=True)
else:
m2 = compose.add_prefix(m1,
prefix,
rename_nodes=rename_nodes,
rename_edges=rename_edges,
rename_inputs=rename_inputs,
rename_outputs=rename_outputs,
rename_initializers=rename_initializers,
rename_value_infos=rename_value_infos)
g_in = m1.graph
g_out = m2.graph
if rename_edges or rename_inputs or rename_outputs or rename_initializers or rename_value_infos:
name_mapping = {}
# Rename inputs/outputs/edges. Propagate name changes from and to edges
if rename_edges:
for n in g_in.node:
for e in n.input:
name_mapping[e] = _prefixed(prefix, e)
for e in n.output:
name_mapping[e] = _prefixed(prefix, e)
else:
if rename_inputs:
for elem in g_in.input:
name_mapping[elem.name] = _prefixed(prefix, elem.name)
if rename_outputs:
for elem in g_in.output:
name_mapping[elem.name] = _prefixed(prefix, elem.name)
if rename_initializers:
for init in g_in.initializer:
name_mapping[init.name] = _prefixed(prefix, init.name)
for sparse_init in g_in.sparse_initializer:
name_mapping[sparse_init.values.name] = \
_prefixed(prefix, sparse_init.values.name)
name_mapping[sparse_init.indices.name] = \
_prefixed(prefix, sparse_init.indices.name)
if rename_value_infos:
for value_info in g_in.output:
name_mapping[value_info.name] = _prefixed(
prefix, value_info.name)
for n1, n0 in zip(g_out.node, g_in.node):
for e1, e0 in zip(n1.input, n0.input):
self.assertEqual(name_mapping.get(e0, e0), e1)
for e1, e0 in zip(n1.output, n0.output):
self.assertEqual(name_mapping.get(e0, e0), e1)
for i1, i0 in zip(g_out.input, g_in.input):
self.assertEqual(name_mapping.get(i0.name, i0.name), i1.name)
for o1, o0 in zip(g_out.output, g_in.output):
self.assertEqual(name_mapping.get(o0.name, o0.name), o1.name)
for init1, init0 in zip(g_out.initializer, g_in.initializer):
self.assertEqual(name_mapping.get(init0.name, init0.name),
init1.name)
for sparse_init1, sparse_init0 in zip(g_out.sparse_initializer,
g_in.sparse_initializer):
self.assertEqual(
name_mapping.get(sparse_init0.values.name,
sparse_init0.values.name),
sparse_init1.values.name)
self.assertEqual(
name_mapping.get(sparse_init0.indices.name,
sparse_init0.indices.name),
sparse_init1.indices.name)
for vi1, vi0 in zip(g_out.value_info, g_in.value_info):
self.assertEqual(name_mapping.get(vi0.name, vi0.name),
vi1.name)
if rename_nodes:
for n1, n0 in zip(g_out.node, g_in.node):
self.assertEqual(_prefixed(prefix, n0.name), n1.name)
# 2. Remove their initializers
new_initializers = [
init
for init in model.graph.initializer
if init.name not in nodes_to_remove and init.name not in inputs_to_remove
]
# 3. Remove nodes
new_nodes = [n for n in model.graph.node if n.name not in nodes_to_remove]
# Get Ouput Tensor Types to create ValueInfo for output info
# by running model on dummy input
temp_model = ModelProto()
temp_model.CopyFrom(model)
for i in new_output_names:
op = ValueInfoProto()
op.name = i
temp_model.graph.output.append(op)
onnx.save(temp_model, "__temp.onnx")
sess = onnxruntime.InferenceSession("__temp.onnx")
sess_inps = sess.get_inputs()
input_dict = {}
for i in sess_inps:
shape = fix_shape(i.shape, batch_size)
typ = get_np_type_from_onnxruntime(i.type)
input_dict[i.name] = np.random.rand(*shape).astype(typ)
output_tensors = sess.run(new_output_names, input_dict)
if os.path.exists("__temp.onnx"):
def add_prefix(
model: ModelProto,
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,
rename_functions: Optional[bool] = True,
inplace: Optional[bool] = False,
) -> ModelProto:
"""Adds a prefix to names of elements in a graph: nodes, edges, inputs, outputs,
initializers, sparse initializer, value infos, and local functions.
It can be used as a utility before merging graphs that have overlapping names.
Empty names are not _prefixed.
Arguments:
model (ModelProto): Model
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 nanes
rename_functions (bool): Whether to prefix local function names
inplace (bool): If True, mutates the model directly.
Otherwise, a copy will be created
"""
if type(model) is not ModelProto:
raise ValueError("model argument is not an ONNX model")
if not inplace:
m = ModelProto()
m.CopyFrom(model)
model = m
add_prefix_graph(
model.graph,
prefix,
rename_nodes=rename_nodes,
rename_edges=rename_edges,
rename_inputs=rename_inputs,
rename_outputs=rename_outputs,
rename_initializers=rename_initializers,
rename_value_infos=rename_value_infos,
inplace=True # No need to create a copy, since it's a new model
)
if rename_functions:
f_name_map = {}
for f in model.functions:
new_f_name = prefix + f.name
f_name_map[f.name] = new_f_name
f.name = new_f_name
# Adjust references to local functions in other local function
# definitions
for f in model.functions:
for n in f.node:
if n.op_type in f_name_map:
n.op_type = f_name_map[n.op_type]
# Adjust references to local functions in the graph
for n in model.graph.node:
if n.op_type in f_name_map:
n.op_type = f_name_map[n.op_type]
return model
def merge_models(m1: ModelProto,
m2: ModelProto,
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,
producer_name: Optional[Text] = 'onnx.compose.merge_models',
producer_version: Optional[Text] = "1.0",
domain: Optional[Text] = "",
model_version: Optional[int] = 1) -> ModelProto:
"""Combines two ONNX models into a single one.
The combined model 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 model.
Both models should have the same IR version, and same operator sets imported.
Arguments:
m1 (ModelProto): First model
m2 (ModelProto): Second model
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 m1
prefix2 (string): Optional prefix to be added to all names in m2
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
producer_name (string): Optional producer name for the combined model. Default: 'onnx.compose'
producer_version (string): Optional producer version for the combined model. Default: "1.0"
domain (string): Optional domain of the combined model. Default: ""
model_version (int): Optional version of the graph encoded. Default: 1
"""
if type(m1) is not ModelProto:
raise ValueError("m1 argument is not an ONNX model")
if type(m2) is not ModelProto:
raise ValueError("m2 argument is not an ONNX model")
if m1.ir_version != m2.ir_version:
raise ValueError(
f"IR version mismatch {m1.ir_version} != {m2.ir_version}."
" Both models should have have the same IR version")
ir_version = m1.ir_version
opset_import_map: MutableMapping[Text, int] = {}
opset_imports = \
[entry for entry in m1.opset_import] + \
[entry for entry in m2.opset_import]
for entry in opset_imports:
if entry.domain in opset_import_map:
found_version = opset_import_map[entry.domain]
if entry.version != found_version:
raise ValueError(
"Can't merge two models with different operator set ids for a given domain. "
f"Got: {m1.opset_import} and {m2.opset_import}")
else:
opset_import_map[entry.domain] = entry.version
# Prefixing names in the graph if requested, adjusting io_map accordingly
if prefix1 or prefix2:
if prefix1:
m1_copy = ModelProto()
m1_copy.CopyFrom(m1)
m1 = m1_copy
m1 = add_prefix(m1, prefix=prefix1)
if prefix2:
m2_copy = ModelProto()
m2_copy.CopyFrom(m2)
m2 = m2_copy
m2 = add_prefix(m2, 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]
graph = merge_graphs(m1.graph,
m2.graph,
io_map,
inputs=inputs,
outputs=outputs,
name=name,
doc_string=doc_string)
model = helper.make_model(graph,
producer_name=producer_name,
producer_version=producer_version,
domain=domain,
model_version=model_version,
opset_imports=opset_imports,
ir_version=ir_version)
# Merging model metadata props
model_props = {}
for meta_entry in m1.metadata_props:
model_props[meta_entry.key] = meta_entry.value
for meta_entry in m2.metadata_props:
if meta_entry.key in model_props:
value = model_props[meta_entry.key]
if value != meta_entry.value:
raise ValueError(
"Can't merge models with different values for the same model metadata property."
f" Found: property = {meta_entry.key}, with values {value} and {meta_entry.value}."
)
else:
model_props[meta_entry.key] = meta_entry.value
helper.set_model_props(model, model_props)
# Merging functions
function_overlap = list(
set([f.name
for f in m1.functions]) & set([f.name for f in m2.functions]))
if function_overlap:
raise ValueError(
"Can't merge models with overlapping local function names."
" Found in both graphs: " + ', '.join(function_overlap))
model.functions.MergeFrom(m1.functions)
model.functions.MergeFrom(m2.functions)
checker.check_model(model)
return model
