From 5d4ff67c362895a324d502eed14af255502e9f80 Mon Sep 17 00:00:00 2001 From: zhangyaobit <1034716+zhangyaobit@users.noreply.github.com> Date: Thu, 17 Mar 2022 20:57:43 -0700 Subject: [PATCH] Support fusion options for benchmark.py (#10900) * Support fusion options for benchmark.py * Add fusion options for tf model export as well. * Add command example and warning related to fusion options. --- .../python/tools/transformers/benchmark.py | 18 ++++++++++++++---- .../python/tools/transformers/onnx_exporter.py | 18 ++++++++++-------- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py index 57741f6e1d..52208ef7ee 100644 --- a/onnxruntime/python/tools/transformers/benchmark.py +++ b/onnxruntime/python/tools/transformers/benchmark.py @@ -34,6 +34,8 @@ python benchmark.py -e torchscript -g -p "fp16" Run ONNXRuntime and TorchScript on CPU for all models with quantization: python benchmark.py -e torchscript onnxruntime -p "int8" -o + Run OnnxRuntime with the ROCM provider and graph optimization script: + python benchmark.py -g -m bert-base-cased --provider rocm --optimizer_info by_script --disable_embed_layer_norm It is recommended to use run_benchmark.sh to launch benchmark. """ @@ -51,6 +53,7 @@ from enum import Enum from benchmark_helper import (OptimizerInfo, create_onnxruntime_session, Precision, setup_logger, get_latency_result, output_details, output_summary, output_fusion_statistics, inference_ort, inference_ort_with_io_binding, allocateOutputBuffers, ConfigModifier) +from fusion_options import FusionOptions from quantize_helper import QuantizeHelper from onnx_exporter import create_onnxruntime_input, load_pretrained_model, export_onnx_model_from_pt, export_onnx_model_from_tf @@ -71,7 +74,7 @@ from transformers import (AutoConfig, AutoTokenizer, AutoModel, GPT2Model, Lxmer def run_onnxruntime(use_gpu, provider, model_names, model_class, config_modifier, precision, num_threads, batch_sizes, sequence_lengths, repeat_times, input_counts, optimizer_info, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, disable_ort_io_binding, use_raw_attention_mask, model_fusion_statistics, - model_source): + model_source, args): import onnxruntime results = [] @@ -92,6 +95,9 @@ def run_onnxruntime(use_gpu, provider, model_names, model_class, config_modifier ) return results + if optimizer_info == OptimizerInfo.NOOPT: + logger.warning(f"OptimizerInfo is set to {optimizer_info}, graph optimizations specified in FusionOptions are not applied.") + for model_name in model_names: all_input_names = MODELS[model_name][0] for num_inputs in input_counts: @@ -99,18 +105,20 @@ def run_onnxruntime(use_gpu, provider, model_names, model_class, config_modifier break input_names = all_input_names[:num_inputs] + args.model_type = MODELS[model_name][3] + fusion_options = FusionOptions.parse(args) if 'pt' in model_source: with torch.no_grad(): onnx_model_file, is_valid_onnx_model, vocab_size, max_sequence_length = export_onnx_model_from_pt( model_name, MODELS[model_name][1], MODELS[model_name][2], MODELS[model_name][3], model_class, config_modifier, cache_dir, onnx_dir, input_names, use_gpu, precision, optimizer_info, - validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics) + validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics, fusion_options) if 'tf' in model_source: onnx_model_file, is_valid_onnx_model, vocab_size, max_sequence_length = export_onnx_model_from_tf( model_name, MODELS[model_name][1], MODELS[model_name][2], MODELS[model_name][3], model_class, config_modifier, cache_dir, onnx_dir, input_names, use_gpu, precision, optimizer_info, - validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics) + validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics, fusion_options) if not is_valid_onnx_model: continue @@ -517,6 +525,8 @@ def parse_arguments(): default=None, help="Manually set the model's layer number") + FusionOptions.add_arguments(parser) + args = parser.parse_args() return args @@ -584,7 +594,7 @@ def main(): args.test_times, args.input_counts, args.optimizer_info, args.validate_onnx, args.cache_dir, args.onnx_dir, args.verbose, args.overwrite, args.disable_ort_io_binding, use_raw_attention_mask, model_fusion_statistics, - args.model_source) + args.model_source, args) except: logger.error(f"Exception", exc_info=True) diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 30d767e930..972b738a40 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -184,13 +184,14 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri def optimize_onnx_model(model_name, onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu, precision, use_raw_attention_mask, overwrite, model_fusion_statistics, - use_external_data_format): + use_external_data_format, optimization_options=None): if overwrite or not os.path.exists(optimized_model_path): Path(optimized_model_path).parent.mkdir(parents=True, exist_ok=True) from optimizer import optimize_model from fusion_options import FusionOptions - optimization_options = FusionOptions(model_type) + if optimization_options == None: + optimization_options = FusionOptions(model_type) optimization_options.use_raw_attention_mask(use_raw_attention_mask) if Precision.FLOAT16 == precision: optimization_options.enable_gelu_approximation = True @@ -317,7 +318,8 @@ def validate_and_optimize_onnx(model_name, onnx_model_path, example_inputs, example_outputs_flatten, - output_names=None): + output_names, + fusion_options): is_valid_onnx_model = True if validate_onnx: is_valid_onnx_model = validate_onnx_model(onnx_model_path, example_inputs, example_outputs_flatten, use_gpu, @@ -330,7 +332,7 @@ def validate_and_optimize_onnx(model_name, False, use_external_data_format) optimize_onnx_model(model_name, onnx_model_path, optimized_model_path, model_type, config.num_attention_heads, config.hidden_size, use_gpu, precision, use_raw_attention_mask, overwrite, - model_fusion_statistics, use_external_data_format) + model_fusion_statistics, use_external_data_format, fusion_options) onnx_model_path = optimized_model_path if validate_onnx: @@ -352,7 +354,7 @@ def validate_and_optimize_onnx(model_name, def export_onnx_model_from_pt(model_name, opset_version, use_external_data_format, model_type, model_class, config_modifier, cache_dir, onnx_dir, input_names, use_gpu, precision, optimizer_info, - validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics): + validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics, fusion_options): config, model = load_pt_model(model_name, model_class, cache_dir, config_modifier) # config, model = load_pt_model_from_tf(model_name) @@ -401,14 +403,14 @@ def export_onnx_model_from_pt(model_name, opset_version, use_external_data_forma onnx_model_file, is_valid_onnx_model, vocab_size = validate_and_optimize_onnx( model_name, use_external_data_format, model_type, onnx_dir, input_names, use_gpu, precision, optimizer_info, validate_onnx, use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, - example_inputs, example_outputs_flatten, None) + example_inputs, example_outputs_flatten, None, fusion_options) return onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size def export_onnx_model_from_tf(model_name, opset_version, use_external_data_format, model_type, model_class, config_modifier, cache_dir, onnx_dir, input_names, use_gpu, precision, optimizer_info, - validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics): + validate_onnx, use_raw_attention_mask, overwrite, model_fusion_statistics, fusion_options): # Use CPU to export import tensorflow as tf tf.config.set_visible_devices([], 'GPU') @@ -495,6 +497,6 @@ def export_onnx_model_from_tf(model_name, opset_version, use_external_data_forma opt_onnx_model_file, onnx_model_file, is_valid_onnx_model, vocab_size = validate_and_optimize_onnx( model_name, use_external_data_format, model_type, onnx_dir, input_names, use_gpu, precision, optimizer_info, validate_onnx, use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, - example_inputs, example_outputs_flatten, output_names) + example_inputs, example_outputs_flatten, output_names, fusion_options) return opt_onnx_model_file, onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size