onnx_light.onnx.utils#

class onnx_light.onnx.utils.Extractor(model: ModelProto)#

Extracts a sub-model from an ONNX model by specifying input and output tensor names.

extract_model(input_names: list[str], output_names: list[str]) → ModelProto#

Extracts a sub-model defined by input_names and output_names.

Returns:

A new ModelProto containing only the nodes reachable from output_names given input_names as graph inputs.

onnx_light.onnx.utils.MAXIMUM_PROTOBUF = 2147483648#

Maximum protobuf size in bytes (2 GB).

onnx_light.onnx.utils.extract_model(input_path: str | PathLike, output_path: str | PathLike, input_names: list[str], output_names: list[str], check_model: bool = True, infer_shapes: bool = False) → None#

Extracts sub-model from an ONNX model.

The sub-model is defined by the names of the input and output tensors exactly.

Note: For control-flow operators, e.g. If and Loop, the boundary of sub-model, which is defined by the input and output tensors, should not cut through the subgraph that is connected to the main graph as attributes of these operators.

Parameters:

  • input_path – The path to original ONNX model.

  • output_path – The path to save the extracted ONNX model.

  • input_names – The names of the input tensors that to be extracted.

  • output_names – The names of the output tensors that to be extracted.

  • check_model – Whether to run model checker on the original model and the extracted model.

  • infer_shapes – Whether to infer the shapes of the original model. Not yet supported; raises NotImplementedError when True.