From 94c3e2048bdeca11c83f85e0164da63f4187c884 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Thu, 19 Aug 2021 13:02:28 -0700 Subject: [PATCH] [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. --- .../python/util/convert_onnx_models_to_ort.py | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/tools/python/util/convert_onnx_models_to_ort.py b/tools/python/util/convert_onnx_models_to_ort.py index d9e800d7fe..5f2a3edb9e 100644 --- a/tools/python/util/convert_onnx_models_to_ort.py +++ b/tools/python/util/convert_onnx_models_to_ort.py @@ -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)