Source code for onnx.utils
# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
import tarfile
from collections import deque
import onnx.checker
import onnx.helper
import onnx.shape_inference
from onnx import FunctionProto, ModelProto, NodeProto, TensorProto, ValueInfoProto
[docs]
class Extractor:
def __init__(self, model: ModelProto) -> None:
self.model = model
self.graph = self.model.graph
self.initializers: dict[str, TensorProto] = self._build_name2obj_dict(
self.graph.initializer
)
self.value_infos: dict[str, ValueInfoProto] = self._build_name2obj_dict(
self.graph.value_info
)
# Add input and output values (not included in the value_info for intermediate values)
self.value_infos.update(self._build_name2obj_dict(self.graph.input))
self.value_infos.update(self._build_name2obj_dict(self.graph.output))
self.outmap: dict[str, int] = self._build_output_dict(self.graph)
@staticmethod
def _build_name2obj_dict(objs) -> dict:
return {obj.name: obj for obj in objs}
@staticmethod
def _build_output_dict(graph) -> dict[str, int]:
output_to_index: dict[str, int] = {}
for index, node in enumerate(graph.node):
for output_name in node.output:
if output_name == "":
continue
assert output_name not in output_to_index # output_name is unique
output_to_index[output_name] = index
return output_to_index
def _collect_new_io(self, io_names_to_extract: list[str]) -> list[ValueInfoProto]:
# Validate that all names exist in self.value_infos
missing_names = [
name for name in io_names_to_extract if name not in self.value_infos
]
if missing_names:
raise ValueError(
f"The following names were not found in value_infos: {', '.join(missing_names)}"
)
return [self.value_infos[name] for name in io_names_to_extract]
def _dfs_search_reachable_nodes(
self,
node_output_name: str,
graph_input_names: set[str],
reachable: set[int],
) -> None:
"""Helper function to find nodes which are connected to an output
Arguments:
node_output_name (str): The name of the output
graph_input_names (set of string): The names of all inputs of the graph
reachable (set of int): The set of indexes to reachable nodes in `nodes`
"""
stack = [node_output_name]
while stack:
current_output_name = stack.pop()
# finish search at inputs
if current_output_name in graph_input_names:
continue
# find nodes connected to this output
if current_output_name in self.outmap:
index = self.outmap[current_output_name]
if index not in reachable:
# add nodes connected to this output to sets
reachable.add(index)
stack += [
input_name
for input_name in self.graph.node[index].input
if input_name != ""
]
def _collect_reachable_nodes(
self,
input_names: list[str],
output_names: list[str],
) -> list[NodeProto]:
_input_names = set(input_names)
reachable: set[int] = set()
for name in output_names:
self._dfs_search_reachable_nodes(name, _input_names, reachable)
# needs to be topologically sorted
return [self.graph.node[index] for index in sorted(reachable)]
def _collect_referred_local_functions(
self,
nodes: list[NodeProto],
) -> list[FunctionProto]:
# a node in a model graph may refer a function.
# a function contains nodes, some of which may in turn refer a function.
# we need to find functions referred by graph nodes and
# by nodes used to define functions.
function_map: dict[tuple[str, str], FunctionProto] = {}
for function in self.model.functions:
function_map[(function.name, function.domain)] = function
referred_local_functions: list[FunctionProto] = []
queue = deque(nodes)
while queue:
node = queue.popleft()
# check if the node is a function op
if (node.op_type, node.domain) in function_map:
function = function_map.pop((node.op_type, node.domain))
referred_local_functions.append(function)
queue.extend(function.node)
# needs to be topologically sorted
return referred_local_functions
def _collect_reachable_tensors(
self,
nodes: list[NodeProto],
) -> tuple[list[TensorProto], list[ValueInfoProto]]:
all_tensors_names: set[str] = set()
for node in nodes:
all_tensors_names.update(node.input)
all_tensors_names.update(node.output)
initializer = [
self.initializers[t] for t in self.initializers if t in all_tensors_names
]
value_info = [
self.value_infos[t] for t in self.value_infos if t in all_tensors_names
]
len_sparse_initializer = len(self.graph.sparse_initializer)
if len_sparse_initializer != 0:
raise ValueError(
f"len_sparse_initializer is {len_sparse_initializer}, it must be 0."
)
len_quantization_annotation = len(self.graph.quantization_annotation)
if len_quantization_annotation != 0:
raise ValueError(
f"len_quantization_annotation is {len_quantization_annotation}, it must be 0."
)
return initializer, value_info
def _make_model(
self,
nodes: list[NodeProto],
inputs: list[ValueInfoProto],
outputs: list[ValueInfoProto],
initializer: list[TensorProto],
value_info: list[ValueInfoProto],
local_functions: list[FunctionProto],
) -> ModelProto:
name = "Extracted from {" + self.graph.name + "}"
graph = onnx.helper.make_graph(
nodes, name, inputs, outputs, initializer=initializer, value_info=value_info
)
meta = {
"ir_version": self.model.ir_version,
"opset_imports": self.model.opset_import,
"producer_name": "onnx.utils.extract_model",
"functions": local_functions,
}
return onnx.helper.make_model(graph, **meta)
def extract_model(
self,
input_names: list[str],
output_names: list[str],
) -> ModelProto:
inputs = self._collect_new_io(input_names)
outputs = self._collect_new_io(output_names)
nodes = self._collect_reachable_nodes(input_names, output_names)
initializer, value_info = self._collect_reachable_tensors(nodes)
local_functions = self._collect_referred_local_functions(nodes)
model = self._make_model(
nodes, inputs, outputs, initializer, value_info, local_functions
)
return model
[docs]
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:
"""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".
Arguments:
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.
"""
if not os.path.exists(input_path):
raise ValueError(f"Invalid input model path: {input_path}")
if not output_path:
raise ValueError("Output model path shall not be empty!")
if not input_names:
raise ValueError("Input tensor names shall not be empty!")
if not output_names:
raise ValueError("Output tensor names shall not be empty!")
if len(input_names) != len(set(input_names)):
raise ValueError("Duplicate names found in the input tensor names.")
if len(output_names) != len(set(output_names)):
raise ValueError("Duplicate names found in the output tensor names.")
if check_model:
onnx.checker.check_model(input_path)
if infer_shapes and os.path.getsize(input_path) > onnx.checker.MAXIMUM_PROTOBUF:
onnx.shape_inference.infer_shapes_path(input_path, output_path)
model = onnx.load(output_path)
elif infer_shapes:
model = onnx.load(input_path, load_external_data=False)
model = onnx.shape_inference.infer_shapes(model)
base_dir = os.path.dirname(input_path)
onnx.load_external_data_for_model(model, base_dir)
else:
model = onnx.load(input_path)
e = Extractor(model)
extracted = e.extract_model(input_names, output_names)
if extracted.ByteSize() > onnx.checker.MAXIMUM_PROTOBUF:
location = os.path.basename(output_path) + ".data"
onnx.save(extracted, output_path, save_as_external_data=True, location=location)
else:
onnx.save(extracted, output_path)
if check_model:
onnx.checker.check_model(output_path)
def _tar_members_filter(
tar: tarfile.TarFile, base: str | os.PathLike
) -> list[tarfile.TarInfo]:
"""Check that the content of ``tar`` will be extracted safely
Args:
tar: The tarball file
base: The directory where the tarball will be extracted
Returns:
list of tarball members
"""
result = []
for member in tar:
member_path = os.path.join(base, member.name)
abs_base = os.path.abspath(base)
abs_member = os.path.abspath(member_path)
if not abs_member.startswith(abs_base):
raise RuntimeError(
f"The tarball member {member_path} in downloading model contains "
f"directory traversal sequence which may contain harmful payload."
)
if member.issym() or member.islnk():
raise RuntimeError(
f"The tarball member {member_path} in downloading model contains "
f"symbolic links which may contain harmful payload."
)
result.append(member)
return result
def _extract_model_safe(
model_tar_path: str | os.PathLike, local_model_with_data_dir_path: str | os.PathLike
) -> None:
"""Safely extracts a tar file to a specified directory.
This function ensures that the extraction process mitigates against
directory traversal vulnerabilities by validating or sanitizing paths
within the tar file. It also provides compatibility for different versions
of the tarfile module by checking for the availability of certain attributes
or methods before invoking them.
Args:
model_tar_path: The path to the tar file to be extracted.
local_model_with_data_dir_path: The directory path where the tar file
contents will be extracted to.
"""
with tarfile.open(model_tar_path) as model_with_data_zipped:
# Mitigate tarball directory traversal risks
if hasattr(tarfile, "data_filter"):
model_with_data_zipped.extractall(
path=local_model_with_data_dir_path, filter="data"
)
else:
model_with_data_zipped.extractall(
path=local_model_with_data_dir_path,
members=_tar_members_filter(
model_with_data_zipped, local_model_with_data_dir_path
),
)