mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
[convert_onnx_models_to_ort.py] Add option to specify NNAPI EP partitioning stop ops. (#8668)
Add option to specify NNAPI EP partitioning stop ops from the ORT format model conversion script.
This commit is contained in:
parent
81889a1cf6
commit
94c3e2048b
1 changed files with 22 additions and 5 deletions
|
|
@ -41,7 +41,8 @@ def _create_config_file_from_ort_models(onnx_model_path_or_dir: pathlib.Path, op
|
|||
|
||||
def _create_session_options(optimization_level: ort.GraphOptimizationLevel,
|
||||
output_model_path: pathlib.Path,
|
||||
custom_op_library: pathlib.Path):
|
||||
custom_op_library: pathlib.Path,
|
||||
session_options_config_entries: typing.Dict[str, str]):
|
||||
so = ort.SessionOptions()
|
||||
so.optimized_model_filepath = str(output_model_path)
|
||||
so.graph_optimization_level = optimization_level
|
||||
|
|
@ -49,11 +50,15 @@ def _create_session_options(optimization_level: ort.GraphOptimizationLevel,
|
|||
if custom_op_library:
|
||||
so.register_custom_ops_library(str(custom_op_library))
|
||||
|
||||
for key, value in session_options_config_entries.items():
|
||||
so.add_session_config_entry(key, value)
|
||||
|
||||
return so
|
||||
|
||||
|
||||
def _convert(model_path_or_dir: pathlib.Path, optimization_level_str: str, use_nnapi: bool, use_coreml: bool,
|
||||
custom_op_library: pathlib.Path, create_optimized_onnx_model: bool, allow_conversion_failures: bool):
|
||||
custom_op_library: pathlib.Path, create_optimized_onnx_model: bool, allow_conversion_failures: bool,
|
||||
session_options_config_entries: typing.Dict[str, str]):
|
||||
|
||||
optimization_level = _get_optimization_level(optimization_level_str)
|
||||
|
||||
|
|
@ -103,14 +108,16 @@ def _convert(model_path_or_dir: pathlib.Path, optimization_level_str: str, use_n
|
|||
# Create an ONNX file with the same optimizations that will be used for the ORT format file.
|
||||
# This allows the ONNX equivalent of the ORT format model to be easily viewed in Netron.
|
||||
optimized_target_path = model.with_suffix(".{}.optimized.onnx".format(optimization_level_str))
|
||||
so = _create_session_options(optimization_level, optimized_target_path, custom_op_library)
|
||||
so = _create_session_options(optimization_level, optimized_target_path, custom_op_library,
|
||||
session_options_config_entries)
|
||||
|
||||
print("Saving optimized ONNX model {} to {}".format(model, optimized_target_path))
|
||||
_ = ort.InferenceSession(str(model), sess_options=so, providers=providers,
|
||||
disabled_optimizers=optimizer_filter)
|
||||
|
||||
# Load ONNX model, optimize, and save to ORT format
|
||||
so = _create_session_options(optimization_level, ort_target_path, custom_op_library)
|
||||
so = _create_session_options(optimization_level, ort_target_path, custom_op_library,
|
||||
session_options_config_entries)
|
||||
so.add_session_config_entry('session.save_model_format', 'ORT')
|
||||
|
||||
print("Converting optimized ONNX model {} to ORT format model {}".format(model, ort_target_path))
|
||||
|
|
@ -192,6 +199,11 @@ def parse_args():
|
|||
parser.add_argument('--allow_conversion_failures', action='store_true',
|
||||
help='Whether to proceed after encountering model conversion failures.')
|
||||
|
||||
parser.add_argument('--nnapi_partitioning_stop_ops',
|
||||
help='Specify the list of NNAPI EP partitioning stop ops. '
|
||||
'In particular, specify the value of the "ep.nnapi.partitioning_stop_ops" session '
|
||||
'options config entry.')
|
||||
|
||||
parser.add_argument('model_path_or_dir', type=pathlib.Path,
|
||||
help='Provide path to ONNX model or directory containing ONNX model/s to convert. '
|
||||
'All files with a .onnx extension, including in subdirectories, will be processed.')
|
||||
|
|
@ -217,8 +229,13 @@ def convert_onnx_models_to_ort():
|
|||
if args.use_coreml and 'CoreMLExecutionProvider' not in ort.get_available_providers():
|
||||
raise ValueError('The CoreML Execution Provider was not included in this build of ONNX Runtime.')
|
||||
|
||||
session_options_config_entries = {}
|
||||
|
||||
if args.nnapi_partitioning_stop_ops is not None:
|
||||
session_options_config_entries["ep.nnapi.partitioning_stop_ops"] = args.nnapi_partitioning_stop_ops
|
||||
|
||||
_convert(model_path_or_dir, args.optimization_level, args.use_nnapi, args.use_coreml, custom_op_library,
|
||||
args.save_optimized_onnx_model, args.allow_conversion_failures)
|
||||
args.save_optimized_onnx_model, args.allow_conversion_failures, session_options_config_entries)
|
||||
|
||||
_create_config_file_from_ort_models(model_path_or_dir, args.optimization_level, args.enable_type_reduction)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue