onnx_helpers

Classes

ModelSplit

Configuration about how to split a model.

Functions

update_model_version(β†’Β onnx.ModelProto)

Update an ONNX model with IR and OPSET supported by Model SDK.

remove_infer_shape(model)

Remove shape inference results in a model.

load_model(β†’Β onnx.ModelProto)

Load a model and update its version information.

save_model(model,Β model_fname[,Β save_only])

Save a model to disk.

run_model(β†’Β List[numpy.ndarray])

Run inference on a model.

verify_models_equal(model_0,Β model_1,Β input_names,Β ...)

Run inference on two models and expect identical match numerically.

verify_models_close(model_0,Β model_1,Β input_names,Β ...)

Run inference on two models and expect close match numerically.

extract_model(β†’Β onnx.ModelProto)

Extract from a model.

split_model(model_splits)

Split a model.

verify_split_models(whole_model,Β split_models,Β ...)

Verify split models by comparing final inference outputs of the whole model and the cascaded splits.

merge_model(β†’Β onnx.ModelProto)

Merge two models.

merge_split_model_with_shared_constant(β†’Β onnx.ModelProto)

Merge two models which may have shared constants.

find_node(β†’Β onnx.NodeProto)

Find an operator node by name.

find_node_output(β†’Β onnx.NodeProto)

Find an operator node whose output is the specified output_name.

find_node_input(β†’Β Tuple[onnx.NodeProto,Β int])

Find an operator node whose input is the specified input_name.

make_node(β†’Β onnx.NodeProto)

Make an operator node following ONNX specification.

remove_node(model,Β node[,Β remove_only])

Remove a node in a model.

insert_node(model,Β node,Β new_node[,Β insert_before,Β ...])

Insert a node in a model.

replace_node(model,Β node,Β new_nodes)

Replace a node in a model.

set_attribute_to_node(model,Β node_name,Β attr_name,Β val)

Set attribute of a node in a model.

remove_nodes_by_name_list(model,Β name_list)

Remove multiple nodes by a list of names.

insert_transpose_pair(β†’Β List[str])

Insert a pair of transpose operators before and after a node.

rewrite_matmul_as_conv(β†’Β List[str])

Rewrite a Matmul operator as a Conv operator.

rewrite_matmul_as_einsum(β†’Β List[str])

Rewrite a Matmul operator as an Einsum operator.

insert_slices_after_node(β†’Β List[str])

Insert Slice operators after a node. All Slice operators output the same length on the slicing axis.

insert_concat(β†’Β str)

Insert Concat operator after a group of nodes.

insert_slices_concat_between_nodes(β†’Β str)

Insert Slices + Concat between two nodes.

rewrite_gemm_as_conv(β†’Β str)

Rewrite a Gemm operator as a Conv operator between two nodes.

connect_nodes(model,Β node_pair,Β out_idx,Β in_idx)

Connect two nodes by feeding the output of one node to the input of next node.

remove_output(model[,Β keep_list])

Run outputs of a model.

remove_outputs_by_name_list(model,Β name_list)

Remove multiple outputs by a list of names.

remove_inputs_by_name_list(model,Β name_list)

Remove multiple inputs by a list of names.

add_io(model,Β io_name,Β io_shape,Β io_dir)

Add input or output to a model.

add_input(model,Β input_name,Β input_shape)

Add input to a model.

add_output(model,Β output_name,Β output_shape)

Add output to a model.

get_io_shape(β†’Β Tuple[int])

Get the shape of a tensor.

update_io_shape(model,Β io_name,Β new_shape)

Update I/O shape in a model.

change_node_output(model,Β node_name,Β new_output_name)

Change output of a node in a model.

is_initializer(β†’Β bool)

Check if a name is an initializer.

find_initializer_value(β†’Β numpy.ndarray)

Find value of an initializer in a model.

remove_initializers(model,Β name_list)

Remove initializers from a list.

add_initializer(model,Β initializer_name,Β initializer_value)

Add an initializer to a model.

remove_duplicated_initializer(model)

Remove duplicated initializers in a model.

transpose_reshape_constant(model,Β param_name,Β perm,Β ...)

Perform transpose and/or reshape on an existing constant initializer.

convert_initializer(model,Β init_name,Β convert_type[,Β ...])

Convert an initializer in a model.

uniquify_initializer_name(β†’Β onnx.ModelProto)

Uniquify initializer names in a model by a prefix name.

export_to_json(model_path[,Β json_file_path])

Convert a model to a simpler json readable representation

update_model_node_names(model_path[,Β updated_model_path])

Update a model, with node names

Module Contents

class onnx_helpers.ModelSplit[source]

Configuration about how to split a model.

Parameters:

  • split_name – Name of the extracted model.

  • input_names – List of input names demarcating the beginning of the split.

  • output_names – List of output names demarcating the end of the split.

  • filename – File name of the extracted model to be saved.

  • parent_split – ModelSplit for the parent if the parent is a split. Cannot be None when parent_path is None.

  • parent_path – Path to the parent model to extract from.

split_name: str[source]
input_names: List[str][source]
output_names: List[str][source]
filename: pathlib.Path | None = None[source]
parent_split: ModelSplit | None = None[source]
parent_path: pathlib.Path | None = None[source]
extract_from_parent_model()[source]
run(input_data: List[numpy.ndarray]) → List[numpy.ndarray][source]
onnx_helpers.update_model_version(model: onnx.ModelProto) → onnx.ModelProto[source]

Update an ONNX model with IR and OPSET supported by Model SDK.

Parameters:

model – Loaded ONNX model in onnx.ModelProto representation.

Returns:

Updated ONNX model in onnx.ModelProto representation.

onnx_helpers.remove_infer_shape(model: onnx.ModelProto)[source]

Remove shape inference results in a model.

Parameters:

model – Loaded model in onnx.ModelProto representation.

onnx_helpers.load_model(model_fname: str, load_only: bool = False) → onnx.ModelProto[source]

Load a model and update its version information.

Parameters:

  • model_fname – File name of the model to load from disk.

  • load_only – Boolean flag, if set to False, to simplify the model after loading.

Returns:

Loaded model in onnx.ModelProto representation.

onnx_helpers.save_model(model: onnx.ModelProto, model_fname: str, save_only: bool = False)[source]

Save a model to disk.

Parameters:

  • model – Model to be saved.

  • model_fname – File name to be used to save the model.

  • save_only – Boolean flag, if set to False, to simplify and re-generate shape inference result before saving to disk.

onnx_helpers.run_model(model_name: str, input_names: List[str], input_data: List[numpy.ndarray]) → List[numpy.ndarray][source]

Run inference on a model.

Parameters:

  • model_name – File name of the model saved on disk.

  • input_names – List of input names expected by the model.

  • input_data – List of input data for the model to run inference on.

Returns:

List of outputs from inference result.

onnx_helpers.verify_models_equal(model_0: str, model_1: str, input_names: List[str], input_data: List[numpy.ndarray])[source]

Run inference on two models and expect identical match numerically.

Parameters:

  • model_0 – File name of the first model.

  • model_1 – File name of the second model.

  • input_names – List of input names expected by both models.

  • input_data – List of input data to run inference on.

onnx_helpers.verify_models_close(model_0: str, model_1: str, input_names: List[str], input_data: List[numpy.ndarray], atol: float = 1e-06)[source]

Run inference on two models and expect close match numerically.

Parameters:

  • model_0 – File name of the first model.

  • model_1 – File name of the second model.

  • input_names – List of input names expected by both models.

  • input_data – List of input data to run inference on.

  • atol – Float value as absolute tolerance in comparison.

onnx_helpers.extract_model(model: onnx.ModelProto, input_names: List[str], output_names: List[str]) → onnx.ModelProto[source]

Extract from a model.

Parameters:

  • model – The model to extract from.

  • input_names – List of input names demarcating the beginning of the split.

  • output_names – List of output names demarcating the end of the split.

Returns:

Extracted model in onnx.ModelProto representation.

onnx_helpers.split_model(model_splits: List[ModelSplit])[source]

Split a model.

Parameters:

model_splits – List of ModelSplit to perform model extraction.

onnx_helpers.verify_split_models(whole_model: str, split_models: List[ModelSplit], input_names: List[str], input_data: List[numpy.ndarray])[source]
Verify split models by comparing final inference outputs of the whole model and the cascaded splits.

The order of the splits must be in the same execution order of the original whole model.

Parameters:

  • whole_model – File name of the original whole model.

  • split_models – List of ModelSplit extracted from the whole model.

  • input_names – List of input names expected by the whole model.

  • input_data – List of input data to run inference on.

onnx_helpers.merge_model(model_0: str | pathlib.Path | onnx.ModelProto, model_1: str | pathlib.Path | onnx.ModelProto, io_map: List[Tuple[str, str]]) → onnx.ModelProto[source]

Merge two models.

Parameters:

  • model_0 – First model provided by file name or loaded as onnx.ModelProto.

  • model_1 – Second model provided by file name or loaded as onnx.ModelProto.

  • io_map – List of tuples mapping output of the first model to input of the second model.

Returns:

Merged model in onnx.ModelProto representation.

onnx_helpers.merge_split_model_with_shared_constant(base_model: onnx.ModelProto | None, model_split: str, constant_prefix: str = None) → onnx.ModelProto[source]
Merge two models which may have shared constants.

When a model is split into multiple parts, some constants are duplicated in multiple parts. Later, when they are merged (possibly after modifications), there are conflicts because of the same constant names. A solution is to add prefix to constant names in the second model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • model_split – File name of the second model to merge.

  • constant_prefix – Prefix name to prepend to all constants in the second model.

Returns:

Merged model in onnx.ModelProto representation.

onnx_helpers.find_node(model: onnx.ModelProto, node_name: str) → onnx.NodeProto[source]

Find an operator node by name.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_name – Name of a node.

Returns:

Operator node in onnx.NodeProto representation.

onnx_helpers.find_node_output(model: onnx.ModelProto, output_name: str) → onnx.NodeProto[source]

Find an operator node whose output is the specified output_name.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • output_name – Name of the output.

Returns:

Operator node in onnx.NodeProto representation.

onnx_helpers.find_node_input(model: onnx.ModelProto, input_name: str) → Tuple[onnx.NodeProto, int][source]

Find an operator node whose input is the specified input_name.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • input_name – Name of the input.

Returns:

Tuple of Operator node and input index.

onnx_helpers.make_node(**kwargs: Any) → onnx.NodeProto[source]

Make an operator node following ONNX specification.

Parameters:

**kwargs –

Operator dependent properties.

Returns:

Newly constructed node in onnx.NodeProto representation.

onnx_helpers.remove_node(model: onnx.ModelProto, node: str | onnx.NodeProto, remove_only: bool = False)[source]

Remove a node in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node – Node to be removed, provided as node name or onnx.NodeProto representation.

  • remove_only – Boolean flag, if set to False, to reconnect surrounding nodes.

onnx_helpers.insert_node(model: onnx.ModelProto, node: onnx.NodeProto, new_node: onnx.NodeProto, insert_before: bool = False, insert_only: bool = False)[source]

Insert a node in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node – Reference node where to insert a new node.

  • new_node – New node to be inserted.

  • insert_before – Boolean flag, True to insert before or False to insert after the reference node.

  • insert_only – Boolean flag, if False, to connect the new node to the reference node.

onnx_helpers.replace_node(model: onnx.ModelProto, node: onnx.NodeProto, new_nodes: List[onnx.NodeProto])[source]

Replace a node in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node – Reference node to replace.

  • new_nodes – List of nodes to replace with.

onnx_helpers.set_attribute_to_node(model: onnx.ModelProto, node_name: str, attr_name: str, val: Any)[source]

Set attribute of a node in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_name – Name of the node.

  • attr_name – Name of the attribute.

  • val – Value of the attribute.

onnx_helpers.remove_nodes_by_name_list(model: onnx.ModelProto, name_list: List[str])[source]

Remove multiple nodes by a list of names.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name_list – List of names of the nodes to be removed.

onnx_helpers.insert_transpose_pair(model: onnx.ModelProto, node_name: str, perm_before: Tuple[int], perm_after: Tuple[int]) → List[str][source]

Insert a pair of transpose operators before and after a node.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_name – Name of the reference node.

  • perm_before – Permutation before the node.

  • perm_after – Permutation after the node.

Returns:

List of names of the newly inserted transpose nodes.

onnx_helpers.rewrite_matmul_as_conv(model: onnx.ModelProto, name_list: List[str]) → List[str][source]

Rewrite a Matmul operator as a Conv operator.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name_list – List of Matmul node names to be rewritten.

Returns:

List of names of the new convolution nodes.

onnx_helpers.rewrite_matmul_as_einsum(model: onnx.ModelProto, eqn_list: Dict[str, str]) → List[str][source]

Rewrite a Matmul operator as an Einsum operator.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • eqn_list – Dictionary for rewrite with key being node name and value being equation string for Einsum.

Returns:

List of names of the new Einsum nodes.

onnx_helpers.insert_slices_after_node(model: onnx.ModelProto, node_name: str, *, axis: int, nslices: int, slice_size: int) → List[str][source]

Insert Slice operators after a node. All Slice operators output the same length on the slicing axis.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_name – Name of the node after which to insert slices.

  • axis – Axis for the Slice operators.

  • nslices – Number of Slice operators.

  • slice_size – Length for the Slice operators on the slicing axis.

Returns:

List of names of the new Slice nodes.

onnx_helpers.insert_concat(model: onnx.ModelProto, input_nodes: List[str], at_node: str, *, axis: int) → str[source]

Insert Concat operator after a group of nodes.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • input_nodes – List of names representing the group of nodes which are inputs to Concat.

  • at_node – Name of node where insertion happens.

  • axis – Axis on which to concatenate.

Returns:

Name of newly inserted Concat node.

onnx_helpers.insert_slices_concat_between_nodes(model: onnx.ModelProto, after_node: str, before_node: str, *, slice_axis: int, nslices: int, slice_size: int, concat_axis: int) → str[source]

Insert Slices + Concat between two nodes.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • after_node – Name of the node after which a group of Slice operators are inserted.

  • before_node – Name of the node before which a Concat operator is inserted.

  • slice_axis – Axis for the Slice operators.

  • nslices – Number of Slice operators.

  • slice_size – Length for the Slice operators on the slicing axis.

  • concat_axis – Axis on which to concatenate.

Returns:

Name of newly inserted Concat node.

onnx_helpers.rewrite_gemm_as_conv(model: onnx.ModelProto, after_node: str, at_node: str, *, w_r: int, w_s: int, w_c: int, w_k: int) → str[source]

Rewrite a Gemm operator as a Conv operator between two nodes.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • after_node – Name of the node after which the Conv operator is inserted.

  • at_node – Name of the node before which the Conv operator is inserted.

  • w_k (w_r, w_s, w_c,) – Shape of the convolution weight as RSCK

Returns:

Name of newly inserted Conv node.

onnx_helpers.connect_nodes(model: onnx.ModelProto, node_pair: List[str], out_idx: int, in_idx: int)[source]

Connect two nodes by feeding the output of one node to the input of next node.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_pair – List of two nodes in execution order.

  • out_idx – The index of output of the first node.

  • in_idx – The index of input of the second node.

onnx_helpers.remove_output(model: onnx.ModelProto, keep_list: List[str] = None)[source]

Run outputs of a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • keep_list – List of output names not to be removed.

onnx_helpers.remove_outputs_by_name_list(model: onnx.ModelProto, name_list: List[str])[source]

Remove multiple outputs by a list of names.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name_list – List of names of the outputs to be removed.

onnx_helpers.remove_inputs_by_name_list(model: onnx.ModelProto, name_list: List[str])[source]

Remove multiple inputs by a list of names.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name_list – List of names of the inputs to be removed.

onnx_helpers.add_io(model: onnx.ModelProto, io_name: str, io_shape: Tuple[int], io_dir: str)[source]

Add input or output to a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • io_name – Name to be added for an input/output.

  • io_shape – Shape of the added input/output.

  • io_dir – Designation of input or output.

onnx_helpers.add_input(model: onnx.ModelProto, input_name: str, input_shape: Tuple[int])[source]

Add input to a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • input_name – Name of the input to be added.

  • input_shape – Shape of the input to be added.

onnx_helpers.add_output(model: onnx.ModelProto, output_name, output_shape)[source]

Add output to a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • output_name – Name of the output to be added.

  • output_shape – Shape of the output to be added.

onnx_helpers.get_io_shape(node: onnx.NodeProto) → Tuple[int][source]

Get the shape of a tensor.

Parameters:

node – Node representing the tensor.

Returns:

Shape of the tensor.

onnx_helpers.update_io_shape(model: onnx.ModelProto, io_name: str, new_shape: Tuple[int])[source]

Update I/O shape in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • io_name – Name of the I/O

  • new_shape – Shape of the I/O

onnx_helpers.change_node_output(model: onnx.ModelProto, node_name: str, new_output_name: str, old_output_name: str = None)[source]

Change output of a node in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • node_name – Name of a node.

  • new_output_name – Name of the new output for the node.

  • old_output_name – Name of previous output to replace.

onnx_helpers.is_initializer(model: onnx.ModelProto, name: str) → bool[source]

Check if a name is an initializer.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name – Name to check on.

Returns:

True if provided name is an initializer.

onnx_helpers.find_initializer_value(model: onnx.ModelProto, initializer_name: str) → numpy.ndarray[source]

Find value of an initializer in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • initializer_name – Name of the initializer.

Returns:

Value of the initializer.

onnx_helpers.remove_initializers(model: onnx.ModelProto, name_list: List[str])[source]

Remove initializers from a list.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • name_list – List of initializer names.

onnx_helpers.add_initializer(model: onnx.ModelProto, initializer_name: str, initializer_value: numpy.ndarray)[source]

Add an initializer to a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • initializer_name – Name of the initializer.

  • initializer_value – Value of the initializer.

onnx_helpers.remove_duplicated_initializer(model: onnx.ModelProto)[source]

Remove duplicated initializers in a model.

Parameters:

model – Loaded model in onnx.ModelProto representation.

onnx_helpers.transpose_reshape_constant(model: onnx.ModelProto, param_name: str, perm: Tuple[int], new_shape: Tuple[int])[source]

Perform transpose and/or reshape on an existing constant initializer.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • param_name – Name of the constant as an initializer.

  • perm – Permutation to be performed.

  • new_shape – New shape to be reshaped.

onnx_helpers.convert_initializer(model: onnx.ModelProto, init_name: str, convert_type: str, convert_func: Callable[[numpy.ndarray], numpy.ndarray] | None = None)[source]

Convert an initializer in a model.

Parameters:

  • model – Loaded model in onnx.ModelProto representation.

  • init_name – Name of the initializer.

  • convert_type – String describing format conversion.

  • convert_func – Custom conversion function for β€œcustom” convert_type.

onnx_helpers.uniquify_initializer_name(prefix_name: str, model: onnx.ModelProto) → onnx.ModelProto[source]

Uniquify initializer names in a model by a prefix name.

Parameters:

  • prefix_name – Name of the prefix to prepend to all initializer names.

  • model – Loaded model in onnx.ModelProto representation.

Returns:

Updated model in onnx.ModelProto representation.

onnx_helpers.export_to_json(model_path: str, json_file_path: str = 'model.json')[source]

Convert a model to a simpler json readable representation :param model: Loaded model in onnx.ModelProto representation. :param the path where to write the json data to

onnx_helpers.update_model_node_names(model_path: str, updated_model_path: str = 'updated_model.onnx')[source]

Update a model, with node names At times, when converting a model from pytorch to onnx the nodes names are not generated. This API helps in updating the nodes with generated names identified by the operator :param model_path: Input modelpath :param updated_model_path: Modified model save path