mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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.
This commit is contained in:
zhangyaobit
GitHub
parent
b86d105153
commit
5d4ff67c36
2 changed files with 24 additions and 12 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue