onnx.utils#

Extractor#

class onnx.utils.Extractor(model: ModelProto)#
__init__(model: ModelProto) → None#
static _build_name2obj_dict(objs)#
_collect_new_inputs(names: List[str]) → List[ValueInfoProto]#
_collect_new_io_core(original_io, io_names_to_extract)#
_collect_new_outputs(names: List[str]) → List[ValueInfoProto]#
_collect_reachable_nodes(input_names: List[str], output_names: List[str]) → List[NodeProto]#
_collect_reachable_tensors(nodes: List[NodeProto]) → Tuple[List[TensorProto], List[ValueInfoProto]]#
_collect_referred_local_functions(nodes: List[NodeProto]) → List[FunctionProto]#
_dfs_search_reachable_nodes(node_output_name: str, graph_input_names: List[str], reachable_nodes: List[NodeProto]) → None#
_make_model(nodes: List[NodeProto], inputs: List[ValueInfoProto], outputs: List[ValueInfoProto], initializer: List[TensorProto], value_info: List[ValueInfoProto], local_functions: List[FunctionProto]) → ModelProto#

extract_model#

onnx.utils.extract_model(input_path: str, output_path: str, input_names: List[str], output_names: List[str], check_model: bool = True) → 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 (string) – The path to original ONNX model.

  • output_path (string) – 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 extracted model.