onnx.utils¶
Extractor¶
- class onnx.utils.Extractor(model: ModelProto)[source]¶
extract_model¶
- 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 = True) None [source]¶
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.
Note: When the extracted model size is larger than 2GB, the extra data will be saved in “output_path.data”.
- Parameters:
input_path (str | os.PathLike) – The path to original ONNX model.
output_path (str | os.PathLike) – The path to save the extracted ONNX model.
input_names (list of string) – The names of the input tensors that to be extracted.
output_names (list of string) – The names of the output tensors that to be extracted.
check_model (bool) – Whether to run model checker on the original model and the extracted model.
infer_shapes (bool) – Whether to infer the shapes of the original model.