codekingpro/portable-devtools
114k
1#!/usr/bin/env python3
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License.
4
5from __future__ import annotations
6
7import argparse
8import contextlib
9import enum
10import os
11import pathlib
12import tempfile
13
14import onnxruntime as ort
15
16from .file_utils import files_from_file_or_dir, path_match_suffix_ignore_case
17from .onnx_model_utils import get_optimization_level
18from .ort_format_model import create_config_from_models
19
20
21class OptimizationStyle(enum.Enum):
22 Fixed = 0
23 Runtime = 1
24
25
26def _optimization_suffix(optimization_level_str: str, optimization_style: OptimizationStyle, suffix: str):
27 return "{}{}{}".format(
28 f".{optimization_level_str}" if optimization_level_str != "all" else "",
29 ".with_runtime_opt" if optimization_style == OptimizationStyle.Runtime else "",
30 suffix,
31 )
32
33
34def _create_config_file_path(
35 model_path_or_dir: pathlib.Path,
36 output_dir: pathlib.Path | None,
37 optimization_level_str: str,
38 optimization_style: OptimizationStyle,
39 enable_type_reduction: bool,
40):
41 config_name = "{}{}".format(
42 "required_operators_and_types" if enable_type_reduction else "required_operators",
43 _optimization_suffix(optimization_level_str, optimization_style, ".config"),
44 )
45
46 if model_path_or_dir.is_dir():
47 return (output_dir or model_path_or_dir) / config_name
48
49 model_config_path = model_path_or_dir.with_suffix(f".{config_name}")
50
51 if output_dir is not None:
52 return output_dir / model_config_path.name
53
54 return model_config_path
55
56
57def _create_session_options(
58 optimization_level: ort.GraphOptimizationLevel,
59 output_model_path: pathlib.Path,
60 custom_op_library: pathlib.Path,
61 session_options_config_entries: dict[str, str],
62):
63 so = ort.SessionOptions()
64 so.optimized_model_filepath = str(output_model_path)
65 so.graph_optimization_level = optimization_level
66
67 if custom_op_library:
68 so.register_custom_ops_library(str(custom_op_library))
69
70 for key, value in session_options_config_entries.items():
71 so.add_session_config_entry(key, value)
72
73 return so
74
75
76def _convert(
77 model_path_or_dir: pathlib.Path,
78 output_dir: pathlib.Path | None,
79 optimization_level_str: str,
80 optimization_style: OptimizationStyle,
81 custom_op_library: pathlib.Path,
82 create_optimized_onnx_model: bool,
83 allow_conversion_failures: bool,
84 target_platform: str,
85 session_options_config_entries: dict[str, str],
86) -> list[pathlib.Path]:
87 model_dir = model_path_or_dir if model_path_or_dir.is_dir() else model_path_or_dir.parent
88 output_dir = output_dir or model_dir
89
90 optimization_level = get_optimization_level(optimization_level_str)
91
92 def is_model_file_to_convert(file_path: pathlib.Path):
93 if not path_match_suffix_ignore_case(file_path, ".onnx"):
94 return False
95 # ignore any files with an extension of .optimized.onnx which are presumably from previous executions
96 # of this script
97 if path_match_suffix_ignore_case(file_path, ".optimized.onnx"):
98 print(f"Ignoring '{file_path}'")
99 return False
100 return True
101
102 models = files_from_file_or_dir(model_path_or_dir, is_model_file_to_convert)
103
104 if len(models) == 0:
105 raise ValueError(f"No model files were found in '{model_path_or_dir}'")
106
107 providers = ["CPUExecutionProvider"]
108
109 # if the optimization level is greater than or equal to 'layout' we manually exclude the NCHWc transformer.
110 # It's not applicable to ARM devices, and creates a device specific model which won't run on all hardware.
111 # If someone really really really wants to run it they could manually create an optimized onnx model first,
112 # or they could comment out this code.
113 optimizer_filter = None
114 if (
115 (optimization_level == ort.GraphOptimizationLevel.ORT_ENABLE_ALL)
116 or (optimization_level == ort.GraphOptimizationLevel.ORT_ENABLE_LAYOUT)
117 ) and target_platform != "amd64":
118 optimizer_filter = ["NchwcTransformer"]
119
120 converted_models = []
121
122 for model in models:
123 try:
124 relative_model_path = model.relative_to(model_dir)
125
126 (output_dir / relative_model_path).parent.mkdir(parents=True, exist_ok=True)
127
128 ort_target_path = (output_dir / relative_model_path).with_suffix(
129 _optimization_suffix(optimization_level_str, optimization_style, ".ort")
130 )
131
132 if create_optimized_onnx_model:
133 # Create an ONNX file with the same optimization level that will be used for the ORT format file.
134 # This allows the ONNX equivalent of the ORT format model to be easily viewed in Netron.
135 # If runtime optimizations are saved in the ORT format model, there may be some difference in the
136 # graphs at runtime between the ORT format model and this saved ONNX model.
137 optimized_target_path = (output_dir / relative_model_path).with_suffix(
138 _optimization_suffix(optimization_level_str, optimization_style, ".optimized.onnx")
139 )
140 so = _create_session_options(
141 optimization_level, optimized_target_path, custom_op_library, session_options_config_entries
142 )
143 if optimization_style == OptimizationStyle.Runtime:
144 # Limit the optimizations to those that can run in a model with runtime optimizations.
145 so.add_session_config_entry("optimization.minimal_build_optimizations", "apply")
146
147 print(f"Saving optimized ONNX model {model} to {optimized_target_path}")
148 _ = ort.InferenceSession(
149 str(model), sess_options=so, providers=providers, disabled_optimizers=optimizer_filter
150 )
151
152 # Load ONNX model, optimize, and save to ORT format
153 so = _create_session_options(
154 optimization_level, ort_target_path, custom_op_library, session_options_config_entries
155 )
156 so.add_session_config_entry("session.save_model_format", "ORT")
157 if optimization_style == OptimizationStyle.Runtime:
158 so.add_session_config_entry("optimization.minimal_build_optimizations", "save")
159
160 print(f"Converting optimized ONNX model {model} to ORT format model {ort_target_path}")
161 _ = ort.InferenceSession(
162 str(model), sess_options=so, providers=providers, disabled_optimizers=optimizer_filter
163 )
164
165 converted_models.append(ort_target_path)
166
167 # orig_size = os.path.getsize(onnx_target_path)
168 # new_size = os.path.getsize(ort_target_path)
169 # print("Serialized {} to {}. Sizes: orig={} new={} diff={} new:old={:.4f}:1.0".format(
170 # onnx_target_path, ort_target_path, orig_size, new_size, new_size - orig_size, new_size / orig_size))
171 except Exception as e:
172 print(f"Error converting {model}: {e}")
173 if not allow_conversion_failures:
174 raise
175
176 print(f"Converted {len(converted_models)}/{len(models)} models successfully.")
177
178 return converted_models
179
180
181def parse_args():
182 parser = argparse.ArgumentParser(
183 os.path.basename(__file__),
184 description="""Convert the ONNX format model/s in the provided directory to ORT format models.
185 All files with a `.onnx` extension will be processed. For each one, an ORT format model will be created in the
186 given output directory, if specified, or the same directory.
187 A configuration file will also be created containing the list of required operators for all
188 converted models. This configuration file should be used as input to the minimal build via the
189 `--include_ops_by_config` parameter.
190 """,
191 )
192
193 parser.add_argument(
194 "--output_dir",
195 type=pathlib.Path,
196 help="Provide an output directory for the converted model/s and configuration file. "
197 "If unspecified, the converted ORT format model/s will be in the same directory as the ONNX model/s.",
198 )
199
200 parser.add_argument(
201 "--optimization_style",
202 nargs="+",
203 default=[OptimizationStyle.Fixed.name, OptimizationStyle.Runtime.name],
204 choices=[e.name for e in OptimizationStyle],
205 help="Style of optimization to perform on the ORT format model. "
206 "Multiple values may be provided. The conversion will run once for each value. "
207 "The general guidance is to use models optimized with "
208 f"'{OptimizationStyle.Runtime.name}' style when using NNAPI or CoreML and "
209 f"'{OptimizationStyle.Fixed.name}' style otherwise. "
210 f"'{OptimizationStyle.Fixed.name}': Run optimizations directly before saving the ORT "
211 "format model. This bakes in any platform-specific optimizations. "
212 f"'{OptimizationStyle.Runtime.name}': Run basic optimizations directly and save certain "
213 "other optimizations to be applied at runtime if possible. This is useful when using a "
214 "compiling EP like NNAPI or CoreML that may run an unknown (at model conversion time) "
215 "number of nodes. The saved optimizations can further optimize nodes not assigned to the "
216 "compiling EP at runtime.",
217 )
218
219 parser.add_argument(
220 "--enable_type_reduction",
221 action="store_true",
222 help="Add operator specific type information to the configuration file to potentially reduce "
223 "the types supported by individual operator implementations.",
224 )
225
226 parser.add_argument(
227 "--custom_op_library",
228 type=pathlib.Path,
229 default=None,
230 help="Provide path to shared library containing custom operator kernels to register.",
231 )
232
233 parser.add_argument(
234 "--save_optimized_onnx_model",
235 action="store_true",
236 help="Save the optimized version of each ONNX model. "
237 "This will have the same level of optimizations applied as the ORT format model.",
238 )
239
240 parser.add_argument(
241 "--allow_conversion_failures",
242 action="store_true",
243 help="Whether to proceed after encountering model conversion failures.",
244 )
245
246 parser.add_argument(
247 "--target_platform",
248 type=str,
249 default=None,
250 choices=["arm", "amd64"],
251 help="Specify the target platform where the exported model will be used. "
252 "This parameter can be used to choose between platform-specific options, "
253 "such as QDQIsInt8Allowed(arm), NCHWc (amd64) and NHWC (arm/amd64) format, different "
254 "optimizer level options, etc.",
255 )
256
257 parser.add_argument(
258 "model_path_or_dir",
259 type=pathlib.Path,
260 help="Provide path to ONNX model or directory containing ONNX model/s to convert. "
261 "All files with a .onnx extension, including those in subdirectories, will be "
262 "processed.",
263 )
264
265 parsed_args = parser.parse_args()
266 parsed_args.optimization_style = [OptimizationStyle[style_str] for style_str in parsed_args.optimization_style]
267 return parsed_args
268
269
270def convert_onnx_models_to_ort(
271 model_path_or_dir: pathlib.Path,
272 output_dir: pathlib.Path | None = None,
273 optimization_styles: list[OptimizationStyle] | None = None,
274 custom_op_library_path: pathlib.Path | None = None,
275 target_platform: str | None = None,
276 save_optimized_onnx_model: bool = False,
277 allow_conversion_failures: bool = False,
278 enable_type_reduction: bool = False,
279):
280 if output_dir is not None:
281 if not output_dir.is_dir():
282 output_dir.mkdir(parents=True)
283 output_dir = output_dir.resolve(strict=True)
284
285 optimization_styles = optimization_styles or []
286
287 # setting optimization level is not expected to be needed by typical users, but it can be set with this
288 # environment variable
289 optimization_level_str = os.getenv("ORT_CONVERT_ONNX_MODELS_TO_ORT_OPTIMIZATION_LEVEL", "all")
290 model_path_or_dir = model_path_or_dir.resolve()
291 custom_op_library = custom_op_library_path.resolve() if custom_op_library_path else None
292
293 if not model_path_or_dir.is_dir() and not model_path_or_dir.is_file():
294 raise FileNotFoundError(f"Model path '{model_path_or_dir}' is not a file or directory.")
295
296 if custom_op_library and not custom_op_library.is_file():
297 raise FileNotFoundError(f"Unable to find custom operator library '{custom_op_library}'")
298
299 session_options_config_entries = {}
300
301 if target_platform is not None and target_platform == "arm":
302 session_options_config_entries["session.qdqisint8allowed"] = "1"
303 else:
304 session_options_config_entries["session.qdqisint8allowed"] = "0"
305
306 for optimization_style in optimization_styles:
307 print(
308 f"Converting models with optimization style '{optimization_style.name}' and level '{optimization_level_str}'"
309 )
310
311 converted_models = _convert(
312 model_path_or_dir=model_path_or_dir,
313 output_dir=output_dir,
314 optimization_level_str=optimization_level_str,
315 optimization_style=optimization_style,
316 custom_op_library=custom_op_library,
317 create_optimized_onnx_model=save_optimized_onnx_model,
318 allow_conversion_failures=allow_conversion_failures,
319 target_platform=target_platform,
320 session_options_config_entries=session_options_config_entries,
321 )
322
323 with contextlib.ExitStack() as context_stack:
324 if optimization_style == OptimizationStyle.Runtime:
325 # Convert models again without runtime optimizations.
326 # Runtime optimizations may not end up being applied, so we need to use both converted models with and
327 # without runtime optimizations to get a complete set of ops that may be needed for the config file.
328 model_dir = model_path_or_dir if model_path_or_dir.is_dir() else model_path_or_dir.parent
329 temp_output_dir = context_stack.enter_context(
330 tempfile.TemporaryDirectory(dir=model_dir, suffix=".without_runtime_opt")
331 )
332 session_options_config_entries_for_second_conversion = session_options_config_entries.copy()
333 # Limit the optimizations to those that can run in a model with runtime optimizations.
334 session_options_config_entries_for_second_conversion["optimization.minimal_build_optimizations"] = (
335 "apply"
336 )
337
338 print(
339 "Converting models again without runtime optimizations to generate a complete config file. "
340 "These converted models are temporary and will be deleted."
341 )
342 converted_models += _convert(
343 model_path_or_dir=model_path_or_dir,
344 output_dir=temp_output_dir,
345 optimization_level_str=optimization_level_str,
346 optimization_style=OptimizationStyle.Fixed,
347 custom_op_library=custom_op_library,
348 create_optimized_onnx_model=False, # not useful as they would be created in a temp directory
349 allow_conversion_failures=allow_conversion_failures,
350 target_platform=target_platform,
351 session_options_config_entries=session_options_config_entries_for_second_conversion,
352 )
353
354 print(
355 f"Generating config file from ORT format models with optimization style '{optimization_style.name}' and level '{optimization_level_str}'"
356 )
357
358 config_file = _create_config_file_path(
359 model_path_or_dir,
360 output_dir,
361 optimization_level_str,
362 optimization_style,
363 enable_type_reduction,
364 )
365
366 create_config_from_models(converted_models, config_file, enable_type_reduction)
367
368
369if __name__ == "__main__":
370 args = parse_args()
371 convert_onnx_models_to_ort(
372 args.model_path_or_dir,
373 output_dir=args.output_dir,
374 optimization_styles=args.optimization_style,
375 custom_op_library_path=args.custom_op_library,
376 target_platform=args.target_platform,
377 save_optimized_onnx_model=args.save_optimized_onnx_model,
378 allow_conversion_failures=args.allow_conversion_failures,
379 enable_type_reduction=args.enable_type_reduction,
380 )
381