[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:
Edward Chen 2021-08-19 13:02:28 -07:00 committed by GitHub
parent 81889a1cf6
commit 94c3e2048b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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)