From 879751f3b7fe7f46b0cb5c1109604338c25ee807 Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Fri, 11 Sep 2020 00:47:37 -0700 Subject: [PATCH] Support Tensorflow benchmarking and onnx export in transformers tool (#5068) * init checkin for tf export and tf benchmark * small fix on argparse * refactor * review comments * review comments --- .../python/tools/transformers/benchmark.py | 133 +++++++++++++-- .../tools/transformers/benchmark_helper.py | 4 +- .../tools/transformers/huggingface_models.py | 17 +- .../tools/transformers/onnx_exporter.py | 155 +++++++++++++----- .../tools/transformers/run_benchmark.sh | 20 ++- 5 files changed, 252 insertions(+), 77 deletions(-) diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py index c1e6fde213..5b92d149e4 100644 --- a/onnxruntime/python/tools/transformers/benchmark.py +++ b/onnxruntime/python/tools/transformers/benchmark.py @@ -52,7 +52,7 @@ from benchmark_helper import (create_onnxruntime_session, Precision, setup_logge output_summary, output_fusion_statistics, inference_ort, inference_ort_with_io_binding, allocateOutputBuffers) from quantize_helper import QuantizeHelper -from onnx_exporter import create_onnxruntime_input, load_pretrained_model, export_onnx_model +from onnx_exporter import create_onnxruntime_input, load_pretrained_model, export_onnx_model_from_pt, export_onnx_model_from_tf logger = logging.getLogger('') @@ -69,7 +69,7 @@ from transformers import (AutoConfig, AutoTokenizer, AutoModel, GPT2Model) def run_onnxruntime(use_gpu, model_names, model_class, precision, batch_sizes, sequence_lengths, repeat_times, input_counts, optimize_onnx, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, - disable_ort_io_binding, use_raw_attention_mask, thread_num, model_fusion_statistics): + disable_ort_io_binding, use_raw_attention_mask, thread_num, model_fusion_statistics, model_source): import onnxruntime results = [] @@ -90,11 +90,18 @@ def run_onnxruntime(use_gpu, model_names, model_class, precision, batch_sizes, s input_names = all_input_names[:num_inputs] - with torch.no_grad(): - onnx_model_file, is_valid_onnx_model, vocab_size, max_sequence_length = export_onnx_model( - model_name, MODELS[model_name][1], MODELS[model_name][2], MODELS[model_name][3], model_class, - cache_dir, onnx_dir, input_names, use_gpu, precision, optimize_onnx, validate_onnx, - use_raw_attention_mask, overwrite, model_fusion_statistics) + 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, cache_dir, + onnx_dir, input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, + model_fusion_statistics) + 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, cache_dir, + onnx_dir, input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, + model_fusion_statistics) + if not is_valid_onnx_model: continue @@ -120,7 +127,8 @@ def run_onnxruntime(use_gpu, model_names, model_class, precision, batch_sizes, s if max_sequence_length is not None and sequence_length > max_sequence_length: continue - ort_inputs = create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names) + input_value_type = numpy.int64 if 'pt' in model_source else numpy.int32 + ort_inputs = create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names, input_value_type) result_template = { "engine": "onnxruntime", @@ -146,9 +154,11 @@ def run_onnxruntime(use_gpu, model_names, model_class, precision, batch_sizes, s model_name, [batch_size, sequence_length])) # Get output sizes from a dummy ort run ort_outputs = ort_session.run(ort_output_names, ort_inputs) + + data_type = numpy.longlong if 'pt' in model_source else numpy.int32 result = inference_ort_with_io_binding(ort_session, ort_inputs, result_template, repeat_times, - ort_output_names, ort_outputs, output_buffers, - max_last_state_size, max_pooler_size, batch_size, device) + ort_output_names, ort_outputs, output_buffers, max_last_state_size, + max_pooler_size, batch_size, device, data_type) logger.info(result) results.append(result) @@ -227,6 +237,84 @@ def run_pytorch(use_gpu, model_names, model_class, precision, batch_sizes, seque return results +def run_tensorflow(use_gpu, model_names, model_class, precision, batch_sizes, sequence_lengths, repeat_times, thread_n, cache_dir, + verbose): + results = [] + + import tensorflow as tf + tf.config.threading.set_intra_op_parallelism_threads(thread_n) + + if not use_gpu: + tf.config.set_visible_devices([], 'GPU') + + if use_gpu and not tf.test.is_built_with_cuda(): + logger.error("Please install Tensorflow-gpu, and use a machine with GPU for testing gpu performance.") + return results + + if use_gpu: # Restrict TensorFlow to only use the first GPU + physical_devices = tf.config.list_physical_devices('GPU') + try: + tf.config.set_visible_devices(physical_devices[0], 'GPU') + except RuntimeError as e: + logger.exception(e) + + if precision == Precision.FLOAT16 or precision == Precision.INT8: + raise NotImplementedError("Mixed precision is currently not supported.") + + for model_name in model_names: + config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) + + model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class, if_tf_model=True) + + tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir) + + max_input_size = tokenizer.max_model_input_sizes[model_name] if model_name in tokenizer.max_model_input_sizes else 1024 + + for batch_size in batch_sizes: + if batch_size <= 0: + continue + + for sequence_length in sequence_lengths: + if max_input_size is not None and sequence_length > max_input_size: + continue + + logger.info("Run Tensorflow on {} with input shape {}".format(model_name, [batch_size, sequence_length])) + + import random + rng = random.Random() + values = [rng.randint(0, config.vocab_size - 1) for i in range(batch_size * sequence_length)] + input_ids = tf.constant(values, shape=(batch_size, sequence_length), dtype=tf.int32) + + try: + model(input_ids, training=False) + + runtimes = timeit.repeat(lambda: model(input_ids, training=False), repeat=repeat_times, number=1) + + result = { + "engine": "tensorflow", + "version": tf.__version__, + "device": "cuda" if use_gpu else "cpu", + "optimizer": "", + "precision": precision, + "io_binding": "", + "model_name": model_name, + "inputs": 1, + "batch_size": batch_size, + "sequence_length": sequence_length, + "datetime": str(datetime.now()), + } + result.update(get_latency_result(runtimes, batch_size)) + logger.info(result) + results.append(result) + except RuntimeError as e: + logger.exception(e) + from numba import cuda + device = cuda.get_current_device() + device.reset() + + return results + + def parse_arguments(): parser = argparse.ArgumentParser() @@ -239,12 +327,20 @@ def parse_arguments(): choices=list(MODELS.keys()), help="Pre-trained models in the list: " + ", ".join(MODELS.keys())) + parser.add_argument("--model_source", + required=False, + nargs=1, + type=str, + default='pt', + choices=['pt', 'tf'], + help="Export onnx from pt or tf") + parser.add_argument('--model_class', required=False, type=str, default=None, - choices=list(MODEL_CLASSES.keys()), - help='Model type selected in the list: ' + ', '.join(MODEL_CLASSES.keys())) + choices=list(MODEL_CLASSES), + help='Model type selected in the list: ' + ', '.join(MODEL_CLASSES)) parser.add_argument("-e", "--engines", @@ -252,7 +348,7 @@ def parse_arguments(): nargs="+", type=str, default=['onnxruntime'], - choices=['onnxruntime', 'torch', 'torchscript'], + choices=['onnxruntime', 'torch', 'torchscript', 'tensorflow'], help="Engines to benchmark") parser.add_argument("-c", @@ -356,10 +452,13 @@ def main(): enable_torch = "torch" in args.engines enable_torchscript = "torchscript" in args.engines enable_onnxruntime = "onnxruntime" in args.engines + enable_tensorflow = "tensorflow" in args.engines results = [] - torch.set_num_threads(cpu_count if args.thread_num <= 0 else args.thread_num) + thread_n = cpu_count if args.thread_num <= 0 else args.thread_num + torch.set_num_threads(thread_n) + logger.debug(torch.__config__.parallel_info()) if enable_torch or enable_torchscript: @@ -374,6 +473,10 @@ def main(): results += run_pytorch(args.use_gpu, args.models, args.model_class, args.precision, args.batch_sizes, args.sequence_lengths, args.test_times, False, args.cache_dir, args.verbose) + if enable_tensorflow: + results += run_tensorflow(args.use_gpu, args.models, args.model_class, args.precision, args.batch_sizes, args.sequence_lengths, + args.test_times, thread_n, args.cache_dir, args.verbose) + model_fusion_statistics = {} if enable_onnxruntime: try: @@ -382,7 +485,7 @@ def main(): args.sequence_lengths, args.test_times, args.input_counts, args.optimize_onnx, args.validate_onnx, args.cache_dir, args.onnx_dir, args.verbose, args.overwrite, args.disable_ort_io_binding, use_raw_attention_mask, args.thread_num, - model_fusion_statistics) + model_fusion_statistics, args.model_source) except: logger.error(f"Exception", exc_info=True) diff --git a/onnxruntime/python/tools/transformers/benchmark_helper.py b/onnxruntime/python/tools/transformers/benchmark_helper.py index 1e1eaf4800..af99daabdc 100644 --- a/onnxruntime/python/tools/transformers/benchmark_helper.py +++ b/onnxruntime/python/tools/transformers/benchmark_helper.py @@ -186,7 +186,7 @@ def inference_ort(ort_session, ort_inputs, result_template, repeat_times, batch_ def inference_ort_with_io_binding(ort_session, ort_inputs, result_template, repeat_times, ort_output_names, ort_outputs, - output_buffers, max_last_state_size, max_pooler_size, batch_size, device): + output_buffers, max_last_state_size, max_pooler_size, batch_size, device, data_type=numpy.longlong): result = {} # Bind inputs and outputs to onnxruntime session @@ -194,7 +194,7 @@ def inference_ort_with_io_binding(ort_session, ort_inputs, result_template, repe # Bind inputs to device for name in ort_inputs.keys(): np_input = torch.from_numpy(ort_inputs[name]).to(device) - io_binding.bind_input(name, np_input.device.type, 0, numpy.longlong, np_input.shape, np_input.data_ptr()) + io_binding.bind_input(name, np_input.device.type, 0, data_type, np_input.shape, np_input.data_ptr()) has_pooler = True if len(ort_output_names) == 2 else False # Bind outputs buffers with the sizes needed if not allocated already if output_buffers["last_state"] is None: diff --git a/onnxruntime/python/tools/transformers/huggingface_models.py b/onnxruntime/python/tools/transformers/huggingface_models.py index 5de382aa08..04c79ed3b6 100644 --- a/onnxruntime/python/tools/transformers/huggingface_models.py +++ b/onnxruntime/python/tools/transformers/huggingface_models.py @@ -4,18 +4,13 @@ # license information. # -------------------------------------------------------------------------- -from transformers import AutoModelForQuestionAnswering -from transformers import AutoModelForSequenceClassification -from transformers import AutoModelWithLMHead -from transformers import AutoModel - # Maps model class name to a tuple of model class -MODEL_CLASSES = { - 'AutoModel': AutoModel, - 'AutoModelWithLMHead': AutoModelWithLMHead, - 'AutoModelForSequenceClassification': AutoModelForSequenceClassification, - 'AutoModelForQuestionAnswering': AutoModelForQuestionAnswering -} +MODEL_CLASSES = [ + 'AutoModel', + 'AutoModelWithLMHead', + 'AutoModelForSequenceClassification', + 'AutoModelForQuestionAnswering' +] # List of pretrained models: https://huggingface.co/transformers/pretrained_models.html # Pretrained model name to a tuple of input names, opset_version, use_external_data_format, optimization model type diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index fa870ddf98..064ce3ccea 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -39,17 +39,17 @@ def restore_torch_functions(): torch.triu = torch_func["triu"] -def create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names): - input_ids = numpy.random.randint(low=0, high=vocab_size - 1, size=(batch_size, sequence_length), dtype=numpy.int64) +def create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names, data_type=numpy.int64): + input_ids = numpy.random.randint(low=0, high=vocab_size - 1, size=(batch_size, sequence_length), dtype=data_type) inputs = {'input_ids': input_ids} if "attention_mask" in input_names: - attention_mask = numpy.ones([batch_size, sequence_length], dtype=numpy.int64) + attention_mask = numpy.ones([batch_size, sequence_length], dtype=data_type) inputs['attention_mask'] = attention_mask if "token_type_ids" in input_names: - segment_ids = numpy.zeros([batch_size, sequence_length], dtype=numpy.int64) + segment_ids = numpy.zeros([batch_size, sequence_length], dtype=data_type) inputs['token_type_ids'] = segment_ids return inputs @@ -95,7 +95,7 @@ def validate_onnx_model(onnx_model_path, example_inputs, example_outputs_flatten logger.info(f"{onnx_model_path} is a valid ONNX model") - # Compare the inference result with PyTorch + # Compare the inference result with PyTorch or Tensorflow example_ort_inputs = {k: t.cpu().numpy() for k, t in example_inputs.items()} example_ort_outputs = test_session.run(None, example_ort_inputs) if len(example_outputs_flatten) != len(example_ort_outputs): @@ -178,6 +178,9 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a optimization_options=optimization_options, use_gpu=use_gpu, only_onnxruntime=False) + if model_type == 'bert_keras': + opt_model.use_dynamic_axes() + model_fusion_statistics[optimized_model_path] = opt_model.get_fused_operator_statistics() if Precision.FLOAT16 == precision: @@ -189,33 +192,78 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a def modelclass_dispatcher(model_name, custom_model_class): if (custom_model_class != None): - return MODEL_CLASSES[custom_model_class] + if (custom_model_class in MODEL_CLASSES): + return custom_model_class + else: + raise Exception("Valid model class: " + ' '.join(MODEL_CLASSES)) if model_name in PRETRAINED_GPT2_MODELS: - return GPT2ModelNoPastState + return "GPT2ModelNoPastState" import re if (re.search('-squad$', model_name) != None): - from transformers import AutoModelForQuestionAnswering - return AutoModelForQuestionAnswering + return "AutoModelForQuestionAnswering" elif (re.search('-mprc$', model_name) != None): - from transformers import AutoModelForSequenceClassification - return AutoModelForSequenceClassification + return "AutoModelForSequenceClassification" elif (re.search('gpt2', model_name) != None): - from transformers import AutoModelWithLMHead - return AutoModelWithLMHead + return "AutoModelWithLMHead" - return AutoModel + return "AutoModel" -def load_pretrained_model(model_name, config, cache_dir, custom_model_class): - model_class = modelclass_dispatcher(model_name, custom_model_class) +def load_pretrained_model(model_name, config, cache_dir, custom_model_class, if_tf_model=False): + if if_tf_model and model_name in PRETRAINED_GPT2_MODELS: + raise NotImplementedError("TFGPT2ModelNoPastState is currently not supported.") + + model_class_name = modelclass_dispatcher(model_name, custom_model_class) + + if if_tf_model: + model_class_name = 'TF' + model_class_name + + transformers_module = __import__("transformers", fromlist=[model_class_name]) + model_class = getattr(transformers_module, model_class_name) + return model_class.from_pretrained(model_name, config=config, cache_dir=cache_dir) -def export_onnx_model(model_name, opset_version, use_external_data_format, model_type, model_class, cache_dir, onnx_dir, - input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, - model_fusion_statistics): +def validate_and_optimize_onnx(model_name, use_external_data_format, model_type, onnx_dir, input_names, use_gpu, + precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, config, + model_fusion_statistics, onnx_model_path, example_inputs, example_outputs_flatten): + 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, + False) + + if optimize_onnx or precision == Precision.FLOAT16 or precision == Precision.INT8: # Use script (optimizer.py) to optimize + optimized_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), True, use_gpu, precision, + False, use_external_data_format) + optimize_onnx_model(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) + + onnx_model_path = optimized_model_path + if validate_onnx: + is_valid_onnx_model = validate_onnx_model(onnx_model_path, example_inputs, example_outputs_flatten, use_gpu, + precision == Precision.FLOAT16) + + if precision == Precision.INT8: + logger.info(f"Quantizing model: {onnx_model_path}") + QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path) + logger.info(f"Finished quantizing model: {onnx_model_path}") + + else: # Use OnnxRuntime to optimize + if is_valid_onnx_model: + ort_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), False, use_gpu, precision, True, + use_external_data_format) + optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite, model_fusion_statistics) + + return onnx_model_path, is_valid_onnx_model, config.vocab_size + + +def export_onnx_model_from_pt(model_name, opset_version, use_external_data_format, model_type, model_class, cache_dir, onnx_dir, + input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, + model_fusion_statistics): + config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) if hasattr(config, 'return_dict'): config.return_dict = False @@ -224,6 +272,9 @@ def export_onnx_model(model_name, opset_version, use_external_data_format, model model.cpu() tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir) + max_input_size = tokenizer.max_model_input_sizes[ + model_name] if model_name in tokenizer.max_model_input_sizes else 1024 + example_inputs = tokenizer.encode_plus("This is a sample input", return_tensors="pt") example_inputs = filter_inputs(example_inputs, input_names) @@ -259,35 +310,53 @@ def export_onnx_model(model_name, opset_version, use_external_data_format, model else: logger.info(f"Skip export since model existed: {onnx_model_path}") - 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, - False) + 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, optimize_onnx, validate_onnx, + use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, example_inputs, example_outputs_flatten) - if optimize_onnx or precision == Precision.FLOAT16 or precision == Precision.INT8: # Use script (optimizer.py) to optimize - optimized_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), True, use_gpu, precision, - False, use_external_data_format) - optimize_onnx_model(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) + return onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size - onnx_model_path = optimized_model_path - if validate_onnx: - is_valid_onnx_model = validate_onnx_model(onnx_model_path, example_inputs, example_outputs_flatten, use_gpu, - precision == Precision.FLOAT16) - if precision == Precision.INT8: - logger.info(f"Quantizing model: {onnx_model_path}") - QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path) - logger.info(f"Finished quantizing model: {onnx_model_path}") +def export_onnx_model_from_tf(model_name, opset_version, use_external_data_format, model_type, model_class, cache_dir, onnx_dir, + input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, + model_fusion_statistics): - else: # Use OnnxRuntime to optimize - if is_valid_onnx_model: - ort_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), False, use_gpu, precision, True, - use_external_data_format) - optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite, model_fusion_statistics) + config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) + model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class, if_tf_model=True) + + model._saved_model_inputs_spec = None + + tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir) max_input_size = tokenizer.max_model_input_sizes[ model_name] if model_name in tokenizer.max_model_input_sizes else 1024 - return onnx_model_path, is_valid_onnx_model, config.vocab_size, max_input_size + example_inputs = tokenizer.encode_plus("This is a sample input", return_tensors="tf", max_length=max_input_size, pad_to_max_length=True, truncation=True) + + example_inputs = filter_inputs(example_inputs, input_names) + + example_outputs = model(example_inputs, training=False) + + # Flatten is needed for gpt2 and distilgpt2. + example_outputs_flatten = flatten(example_outputs) + example_outputs_flatten = update_flatten_list(example_outputs_flatten, []) + + onnx_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), False, use_gpu, precision, False, + use_external_data_format) + + if overwrite or not os.path.exists(onnx_model_path): + logger.info("Exporting ONNX model to {}".format(onnx_model_path)) + import keras2onnx + onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=opset_version) + keras2onnx.save_model(onnx_model, onnx_model_path) + else: + logger.info(f"Skip export since model existed: {onnx_model_path}") + + model_type = model_type + '_keras' + + 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, optimize_onnx, validate_onnx, + use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, example_inputs, example_outputs_flatten) + + return onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size + diff --git a/onnxruntime/python/tools/transformers/run_benchmark.sh b/onnxruntime/python/tools/transformers/run_benchmark.sh index 78b4642235..f826cd5118 100644 --- a/onnxruntime/python/tools/transformers/run_benchmark.sh +++ b/onnxruntime/python/tools/transformers/run_benchmark.sh @@ -20,6 +20,7 @@ run_install=true run_ort=true run_torch=false run_torchscript=true +run_tensorflow=false # Devices to test (You can run either CPU or GPU, but not both: gpu need onnxruntime-gpu, and CPU need onnxruntime). run_gpu_fp32=true @@ -52,7 +53,7 @@ models_to_test="bert-base-cased roberta-base gpt2" # export CUDA_VISIBLE_DEVICES=1 # This script will generate a logs file with a list of commands used in tests. -echo echo "ort=$run_ort torch=$run_torch torchscript=$run_torchscript gpu_fp32=$run_gpu_fp32 gpu_fp16=$run_gpu_fp16 cpu=$run_cpu optimizer=$use_optimizer batch=$batch_sizes sequence=$sequence_length models=$models_to_test" >> benchmark.log +echo echo "ort=$run_ort torch=$run_torch torchscript=$run_torchscript tensorflow=$run_tensorflow gpu_fp32=$run_gpu_fp32 gpu_fp16=$run_gpu_fp16 cpu=$run_cpu optimizer=$use_optimizer batch=$batch_sizes sequence=$sequence_length models=$models_to_test" >> benchmark.log # Set it to false to skip testing. You can use it to dry run this script with the log file. run_tests=true @@ -90,7 +91,7 @@ if [ "$run_install" = true ] ; then fi if [ "$run_cli" = true ] ; then - echo "Use onnxruntime_tools.transformers.benchmark" + echo "Use onnxruntime_tools.transformers.benchmark" benchmark_script="-m onnxruntime_tools.transformers.benchmark" else benchmark_script="benchmark.py" @@ -128,6 +129,13 @@ run_one_test() { python $benchmark_script -e torchscript -m $1 $benchmark_options $2 $3 $4 fi fi + + if [ "$run_tensorflow" = true ] ; then + echo python $benchmark_script -e tensorflow -m $1 $benchmark_options $2 $3 $4 >> benchmark.log + if [ "$run_tests" = true ] ; then + python $benchmark_script -e tensorflow -m $1 $benchmark_options $2 $3 $4 + fi + fi } # ------------------------------------------- @@ -151,9 +159,9 @@ if [ "$run_cpu_fp32" = true ] ; then for m in $models_to_test do echo Run CPU Benchmark on model ${m} - run_one_test "${m}" + run_one_test "${m}" done -fi +fi if [ "$run_cpu_int8" = true ] ; then for m in $models_to_test @@ -161,7 +169,7 @@ if [ "$run_cpu_int8" = true ] ; then echo Run CPU Benchmark on model ${m} run_one_test "${m}" -p int8 done -fi +fi if [ "run_tests" = false ] ; then more $log_file @@ -170,4 +178,4 @@ fi # Remove duplicated lines awk '!x[$0]++' ./result.csv > summary_result.csv awk '!x[$0]++' ./fusion.csv > summary_fusion.csv -awk '!x[$0]++' ./detail.csv > summary_detail.csv +awk '!x[$0]++' ./detail.csv > summary_detail.csv \ No newline at end of file