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

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)
 

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:

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

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
+

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