Implementation:Onnx Onnx Model Subgraph Extractor

Knowledge Sources	Onnx_Onnx
Domains	Model Extraction, Graph Utilities, Model Debugging
Last Updated	2026-02-10 00:00 GMT

Overview

Concrete tool for extracting sub-models from ONNX models by specifying input and output tensor names, along with safe tarball extraction utilities, provided by the ONNX library.

Description

The utils.py module provides the Extractor class and extract_model function for isolating subgraphs from ONNX models. Given a set of input and output tensor names, these utilities identify and extract only the nodes, initializers, value_info entries, and local functions needed to compute the specified outputs from the specified inputs.

The Extractor class is initialized with a ModelProto and builds internal dictionaries that map tensor names to their corresponding objects (initializers, value_infos) and map output names to node indices. The extraction process proceeds in several steps:

1. _collect_new_io validates that the requested input and output names exist in the model's value_info (which includes graph inputs, outputs, and intermediate values).
2. _dfs_search_reachable_nodes performs an iterative depth-first search backward from each requested output tensor, following input edges until reaching the requested input tensors. This identifies the minimal set of nodes required.
3. _collect_reachable_tensors gathers all initializers and value_info entries referenced by the reachable nodes.
4. _collect_referred_local_functions uses a BFS traversal starting from the reachable nodes to discover all model-local functions that are directly or transitively referenced.
5. _make_model assembles a new ModelProto with the extracted subgraph, preserving IR version, opset imports, and producer metadata.

The top-level extract_model function provides a file-based interface that loads a model, optionally runs shape inference (needed to populate value_info for intermediate tensors), extracts the subgraph, and saves the result. It handles models larger than 2GB by saving external data to a companion file.

The module also includes tarball safety utilities. _tar_members_filter validates that tar archive members do not contain directory traversal sequences or symbolic links. _extract_model_safe wraps tarfile extraction with these security checks, using Python's built-in data_filter when available (Python 3.12+) or falling back to the custom filter.

Usage

Use this module to isolate specific parts of large ONNX models for debugging, testing, profiling, or deployment. Common use cases include extracting a backbone network from a larger model, isolating a problematic subgraph for bug reporting, creating smaller test models, and removing preprocessing or postprocessing stages. The tarball extraction utilities are used internally by the model hub for safe model archive handling.

Code Reference

Source Location

• Repository: Onnx_Onnx
• File: onnx/utils.py
• Lines: 1-320

Signature

class Extractor:
    def __init__(self, model: ModelProto) -> None: ...
    def extract_model(
        self,
        input_names: list[str],
        output_names: list[str],
    ) -> ModelProto: ...

def extract_model(
    input_path: str | os.PathLike,
    output_path: str | os.PathLike,
    input_names: list[str],
    output_names: list[str],
    check_model: bool = True,
    infer_shapes: bool = True,
) -> None: ...

Internal helper functions:

def _tar_members_filter(
    tar: tarfile.TarFile, base: str | os.PathLike
) -> list[tarfile.TarInfo]: ...

def _extract_model_safe(
    model_tar_path: str | os.PathLike,
    local_model_with_data_dir_path: str | os.PathLike,
) -> None: ...

Import

from onnx.utils import extract_model, Extractor

I/O Contract

Inputs

Name	Type	Required	Description
model	ModelProto	Yes (Extractor)	The ONNX model from which to extract a sub-model
input_path	str or os.PathLike	Yes (extract_model)	Path to the original ONNX model file
output_path	str or os.PathLike	Yes (extract_model)	Path where the extracted model will be saved
input_names	list[str]	Yes	Names of the input tensors that define the boundary of the sub-model
output_names	list[str]	Yes	Names of the output tensors that define the boundary of the sub-model
check_model	bool	No	Whether to run onnx.checker on the original and extracted models (default: True)
infer_shapes	bool	No	Whether to run shape inference before extraction to populate intermediate value_info (default: True)

Outputs

Name	Type	Description
extracted_model	ModelProto	The extracted sub-model containing only the required nodes, initializers, and functions (Extractor.extract_model)
(saved file)	None	The extract_model function saves the result to output_path and returns None

Usage Examples

import onnx
from onnx.utils import extract_model, Extractor

# File-based extraction (most common usage)
extract_model(
    input_path="full_model.onnx",
    output_path="sub_model.onnx",
    input_names=["input_image"],
    output_names=["features"],
)

# Extract without shape inference or model checking
extract_model(
    input_path="model.onnx",
    output_path="extracted.onnx",
    input_names=["X"],
    output_names=["Y", "Z"],
    check_model=False,
    infer_shapes=False,
)

# Programmatic extraction using the Extractor class
model = onnx.load("model.onnx")
model = onnx.shape_inference.infer_shapes(model)
extractor = Extractor(model)
extracted = extractor.extract_model(
    input_names=["conv1_output"],
    output_names=["pool5_output"],
)
onnx.save(extracted, "backbone.onnx")

# Extract multiple outputs
extracted = extractor.extract_model(
    input_names=["input"],
    output_names=["layer1_out", "layer2_out", "layer3_out"],
)

Related Pages

• Environment:Onnx_Onnx_Python_Runtime_Environment