Spaces:
Sleeping
Sleeping
File size: 8,600 Bytes
dc2106c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
import onnx.checker
import onnx.helper
import onnx.shape_inference
from onnx import FunctionProto, ModelProto, NodeProto, TensorProto, ValueInfoProto
class Extractor:
def __init__(self, model: ModelProto) -> None:
self.model = onnx.shape_inference.infer_shapes(model)
self.graph = self.model.graph
self.wmap = self._build_name2obj_dict(self.graph.initializer)
self.vimap = self._build_name2obj_dict(self.graph.value_info)
@staticmethod
def _build_name2obj_dict(objs): # type: ignore
return {obj.name: obj for obj in objs}
def _collect_new_io_core(self, original_io, io_names_to_extract): # type: ignore
original_io_map = self._build_name2obj_dict(original_io)
original_io_names = set(original_io_map)
s_io_names_to_extract = set(io_names_to_extract)
io_names_to_keep = s_io_names_to_extract & original_io_names
new_io_names_to_add = s_io_names_to_extract - original_io_names
new_io_tensors = []
for name in io_names_to_keep:
new_io_tensors.append(original_io_map[name])
for name in new_io_names_to_add:
# activation become input or output
new_io_tensors.append(self.vimap[name])
# adjust sequence
new_io_tensors_map = self._build_name2obj_dict(new_io_tensors)
return [new_io_tensors_map[name] for name in io_names_to_extract]
def _collect_new_inputs(self, names: list[str]) -> list[ValueInfoProto]:
return self._collect_new_io_core(self.graph.input, names) # type: ignore
def _collect_new_outputs(self, names: list[str]) -> list[ValueInfoProto]:
return self._collect_new_io_core(self.graph.output, names) # type: ignore
def _dfs_search_reachable_nodes(
self,
node_output_name: str,
graph_input_names: list[str],
reachable_nodes: list[NodeProto],
) -> None:
if node_output_name in graph_input_names:
return
for node in self.graph.node:
# check output_name first to reduce run time
if node_output_name not in node.output:
continue
if node in reachable_nodes:
continue
reachable_nodes.append(node)
for name in node.input:
self._dfs_search_reachable_nodes(
name, graph_input_names, reachable_nodes
)
def _collect_reachable_nodes(
self,
input_names: list[str],
output_names: list[str],
) -> list[NodeProto]:
reachable_nodes = [] # type: ignore[var-annotated]
for name in output_names:
self._dfs_search_reachable_nodes(name, input_names, reachable_nodes)
# needs to be topology sorted.
nodes = [n for n in self.graph.node if n in reachable_nodes]
return nodes
def _collect_referred_local_functions(
self,
nodes, # type: list[NodeProto]
): # type: (...) -> 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.
def find_referred_funcs(nodes, referred_local_functions): # type: ignore
new_nodes = [] # type: list[NodeProto]
for node in nodes:
# check if the node is a function op
match_function = next(
(
f
for f in self.model.functions
if f.name == node.op_type and f.domain == node.domain
),
None,
)
if match_function and match_function not in referred_local_functions:
referred_local_functions.append(match_function)
new_nodes.extend(match_function.node)
return new_nodes
referred_local_functions = [] # type: list[FunctionProto]
new_nodes = find_referred_funcs(nodes, referred_local_functions)
while new_nodes:
new_nodes = find_referred_funcs(new_nodes, referred_local_functions)
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.wmap[t] for t in self.wmap if t in all_tensors_names]
value_info = [self.vimap[t] for t in self.vimap 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_inputs(input_names)
outputs = self._collect_new_outputs(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
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,
) -> 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.
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 extracted 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 output_names:
raise ValueError("Output tensor names shall not be empty!")
onnx.checker.check_model(input_path)
model = onnx.load(input_path)
e = Extractor(model)
extracted = e.extract_model(input_names, output_names)
onnx.save(extracted, output_path)
if check_model:
onnx.checker.check_model(output_path)
|