tools.onnx_manipulations#

enumerate_onnx_node_types#

onnx_extended.tools.onnx_tools.enumerate_onnx_node_types(model: str | ModelProto | GraphProto, level: int = 0, shapes: Dict[str, TypeProto] | None = None, external: bool = True) → Generator[Dict[str, str | float], None, None][source]#

Looks into types for every node in a model.

Parameters:

  • model – a string or a proto

  • level – level (recursivity level)

  • shapes – known shapes, returned by :func:onnx.shape_inference.infer_shapes`

  • external – loads the external data if the model is loaded

Returns:

a list of dictionary which can be turned into a dataframe.

get_hidden_inputs#

onnx_extended.tools.onnx_manipulations.get_hidden_inputs(nodes: Iterable[NodeProto]) → Set[str][source]#

Returns the list of hidden inputs used by subgraphs.

Parameters:

nodes – list of nodes

Returns:

list of names

onnx_remove_node_unused#

onnx_extended.tools.onnx_manipulations.onnx_remove_node_unused(onnx_model, recursive=True, debug_info=None, **options)[source]#

Removes unused nodes of the graph. An unused node is not involved in the output computation.

Parameters:

  • onnx_model – onnx model

  • recursive – looks into subgraphs

  • debug_info – debug information (private)

  • options – unused

Returns:

new onnx _model

select_model_inputs_outputs#

onnx_extended.tools.onnx_manipulations.select_model_inputs_outputs(model: ModelProto, outputs: List[str] | None = None, inputs: List[str] | None = None, infer_shapes: bool = True, overwrite: Dict[str, Any] | None = None, remove_unused: bool = True, verbose: int = 0)[source]#

Takes a model and changes its outputs.

Parameters:

  • model – ONNX model

  • inputs – new inputs, same ones if None

  • outputs – new outputs, same ones if None

  • infer_shapes – infer inputs and outputs shapes

  • overwrite – overwrite type and shapes for inputs or outputs, overwrite is a dictionary {‘name’: (numpy dtype, shape)}

  • remove_unused – remove unused nodes from the graph

  • verbose – display information while converting

Returns:

modified model

The function removes unneeded nodes.

The following example shows how to change the inputs of model to bypass the first nodes. Shape inferences fails to determine the new inputs type. They need to be overwritten. verbose=1 shows the number of deleted nodes.

import onnx
from onnx_extended.tools.onnx_manipulations import select_model_inputs_outputs

onx = onnx.load(path)
onx2 = select_model_inputs_outputs(
    onx, inputs=["a", "b"],
    infer_shapes=True, verbose=1,
    overwrite={'a': (numpy.int32, None), 'b': (numpy.int64, None)})
onnx.save(onx2, path2)