utils

Utility functions related to onnx.

Functions

duplicate_shared_constants

Duplicate constant tensors if they are shared.

find_lowest_common_ancestor

Function to find the lowest common ancestor of two nodes.

gen_random_inputs

This function generates random inputs for an onnx model.

get_all_input_names

This function returns the inputs names of the given onnx model.

get_batch_size

Returns the batch size of the given onnx model.

get_batch_size_from_bytes

Returns the batch size of the given onnx model.

get_child_nodes

Returns list of output consumer nodes for the given node.

get_input_names

This function returns the external inputs names of the given onnx model.

get_input_names_from_bytes

This function returns the inputs names of the given onnx model in bytes.

get_input_shapes

This function returns the inputs shapes for the given onnx model.

get_input_shapes_from_bytes

This function returns the input shapes of the given onnx model in bytes.

get_node_names

This function returns all node names from the given onnx model.

get_node_names_from_bytes

This function returns all node names from the given onnx model in bytes.

get_output_names

This function returns the output names of the given onnx model.

get_output_names_from_bytes

This function returns the output names of the given onnx model in bytes.

get_output_shapes

This function returns the output shapes for the given onnx model.

get_parent_nodes

Returns list of input producer nodes for the given node.

get_variable_inputs

Returns the variable inputs of the given Node.

is_valid_onnx_model

Checks if the given file is a valid ONNX model.

name_onnx_nodes

Assigns name to the onnx nodes if not present and return the modified status.

parse_shapes_spec

Parse shapes spec and returns them in a dictionary.

randomize_weights

Assigns random values to the onnx model weights.

randomize_weights_onnx_bytes

Assigns random values to the onnx model weights.

remove_weights_data

Removes raw weight data from the onnx model.

save_onnx

Save an ONNX model to given path.

save_onnx_bytes_to_dir

Saves the onnx bytes to a directory with specified file name.

udpate_domain

Updates the domain of all the nodes of the specified op_type to the specified domain.

validate_batch_size

Returns True if all the model inputs has batch dimension equal to batch_size.

validate_onnx

Returns True if the onnx_bytes is valid, else False.

duplicate_shared_constants(onnx_model)

Duplicate constant tensors if they are shared.

Parameters:

onnx_model (ModelProto) –

Return type:

Tuple[ModelProto, bool]

find_lowest_common_ancestor(node1, node2)

Function to find the lowest common ancestor of two nodes.

Parameters:
  • node1 (Node) – First node name.

  • node2 (Node) – Second node name.

Returns:

LCA node. Distance from first node. Distance from second node.

Return type:

Tuple[str | None, int, int]

gen_random_inputs(model, shapes_spec=None)

This function generates random inputs for an onnx model.

Parameters:
  • model (ModelProto) – Loaded in-memory onnx ModelProto.

  • shapes_spec (str) –

Returns:

Dictionary of numpy tensors.

Return type:

Dict[str, ndarray]

get_all_input_names(model)

This function returns the inputs names of the given onnx model.

Parameters:

model (ModelProto) –

Return type:

List[str]

get_batch_size(model)

Returns the batch size of the given onnx model.

Assertion will fail if batch size is not same over all the inputs.

Parameters:

model (ModelProto) –

Return type:

int

get_batch_size_from_bytes(onnx_bytes)

Returns the batch size of the given onnx model.

Assertion will fail if batch size is not same over all the inputs.

Parameters:

onnx_bytes (bytes) –

Return type:

int

get_child_nodes(node)

Returns list of output consumer nodes for the given node.

Parameters:

node (Node) –

Return type:

List[Node]

get_input_names(model, external_inputs_only=True)

This function returns the external inputs names of the given onnx model.

Note: external_input_names = input_names - initializer_names

Parameters:
  • model (ModelProto) – Loaded in-memory onnx ModelProto.

  • external_inputs_only (bool) –

Returns:

List of external input names of the model.

Return type:

List[str]

get_input_names_from_bytes(model_bytes, external_inputs_only=True)

This function returns the inputs names of the given onnx model in bytes.

Parameters:
  • model_bytes (bytes) – Onnx model in bytes.

  • external_inputs_only (bool) –

Returns:

List of input names of the model.

Return type:

List[str]

get_input_shapes(model, external_inputs_only=True)

This function returns the inputs shapes for the given onnx model.

Parameters:
  • model (ModelProto) –

  • external_inputs_only (bool) –

Return type:

Dict[str, List[int]]

get_input_shapes_from_bytes(model_bytes)

This function returns the input shapes of the given onnx model in bytes.

Parameters:

model_bytes (bytes) – Onnx model in bytes.

Returns:

Dictionary of inputs names and shapes.

Return type:

Dict[str, List[int]]

get_node_names(model)

This function returns all node names from the given onnx model.

Parameters:

model (ModelProto) – Loaded in-memory onnx ModelProto.

Returns:

List of node names of the model.

Return type:

List[str]

get_node_names_from_bytes(model_bytes)

This function returns all node names from the given onnx model in bytes.

Parameters:
  • model – onnx model in bytes.

  • model_bytes (bytes) –

Returns:

List of node names of the model.

Return type:

List[str]

get_output_names(model)

This function returns the output names of the given onnx model.

Parameters:

model (ModelProto) – Loaded in-memory onnx ModelProto.

Returns:

List of output names of the model.

Return type:

List[str]

get_output_names_from_bytes(model_bytes)

This function returns the output names of the given onnx model in bytes.

Parameters:

model_bytes (bytes) – Onnx model in bytes.

Returns:

List of output names of the model.

Return type:

List[str]

get_output_shapes(model)

This function returns the output shapes for the given onnx model.

Parameters:

model (ModelProto) –

Return type:

Dict[str, List[int]]

get_parent_nodes(node)

Returns list of input producer nodes for the given node.

Parameters:

node (Node) –

Return type:

List[Node]

get_variable_inputs(node)

Returns the variable inputs of the given Node.

Parameters:

node (Node) –

Return type:

List[Variable]

is_valid_onnx_model(file_path)

Checks if the given file is a valid ONNX model.

name_onnx_nodes(graph)

Assigns name to the onnx nodes if not present and return the modified status.

Parameters:

graph (GraphProto) –

Return type:

bool

parse_shapes_spec(shapes_spec)

Parse shapes spec and returns them in a dictionary.

Example shapes spec: input0:1x3x256x256,input1:1x3x128x128

Parameters:

shapes_spec (str) –

Return type:

Dict[str, List[int]]

randomize_weights(onnx_path)

Assigns random values to the onnx model weights.

Parameters:

onnx_path (str) –

Return type:

None

randomize_weights_onnx_bytes(onnx_bytes, seed=0)

Assigns random values to the onnx model weights.

Parameters:
  • onnx_bytes (bytes) –

  • seed (int) –

Return type:

bytes

remove_weights_data(onnx_bytes)

Removes raw weight data from the onnx model.

Parameters:

onnx_bytes (bytes) –

Return type:

bytes

save_onnx(model, onnx_path, save_as_external_data=False)

Save an ONNX model to given path. If a model is larger than 2GB, will save with external data.

Parameters:
  • model (ModelProto) –

  • onnx_path (str) –

  • save_as_external_data (bool) –

save_onnx_bytes_to_dir(onnx_bytes, onnx_dir, onnx_name)

Saves the onnx bytes to a directory with specified file name.

Parameters:
  • onnx_bytes (bytes) –

  • onnx_dir (str) –

  • onnx_name (str) –

Return type:

None

udpate_domain(onnx_model, op_type, domain)

Updates the domain of all the nodes of the specified op_type to the specified domain.

Parameters:
  • onnx_model (ModelProto) –

  • op_type (str) –

  • domain (str) –

Return type:

ModelProto

validate_batch_size(onnx_bytes, batch_size)

Returns True if all the model inputs has batch dimension equal to batch_size.

Parameters:
  • onnx_bytes (bytes) –

  • batch_size (int) –

Return type:

bool

validate_onnx(onnx_bytes)

Returns True if the onnx_bytes is valid, else False.

Parameters:

onnx_bytes (bytes) –

Return type:

bool