Bert Optimization Script Improvements (#3387)

Add opt_level option for graph optimization level in bert perf test.
Support BERT models that output each layer, where SkipLayerNormalization has more than 4 children.
Check weight and bias are 1D for layer norm fusion.
Add a dummy class Gpt2OnnxModel for further changes of GPT2 model.

onnxruntime/python/tools/bert/BertOnnxModel.py

@@ -38,9 +38,6 @@ class BertOnnxModel(OnnxModel):
 
         self.bert_inputs = []
 
-    def normalize_children_types(self):
-        return ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']
-
     def cast_graph_input_to_int32(self, input_name):
         graph_input = self.find_graph_input(input_name)
         if graph_input is not None and graph_input.type.tensor_type.elem_type != TensorProto.INT32:
@@ -165,21 +162,26 @@ class BertOnnxModel(OnnxModel):
         for normalize_node in skip_layer_norm_nodes:
             # SkipLayerNormalization has two inputs, and one of them is the
             # root input for attention.
-            qkv_nodes = None
-            root_input = None
+            qkv_nodes = self.match_parent_path(normalize_node, ['Add', 'MatMul', 'Reshape', 'Transpose', 'MatMul'],
+                                               [None, 0, 0, 0, 0])
+            if qkv_nodes is None:
+                continue
+
+            other_inputs = []
             for i, input in enumerate(normalize_node.input):
                 if input not in output_name_to_node:
                     continue
-                children = input_name_to_nodes[input]
-                children_types = sorted([child.op_type for child in children])
-                if children_types != self.normalize_children_types():
-                    qkv_nodes = self.match_parent_path(normalize_node,
-                                                       ['Add', 'MatMul', 'Reshape', 'Transpose', 'MatMul'],
-                                                       [i, 0, 0, 0, 0])
-                else:
-                    root_input = input
 
-            if root_input is None or qkv_nodes is None:
+                if input == qkv_nodes[0].output[0]:
+                    continue
+                other_inputs.append(input)
+            if len(other_inputs) != 1:
+                continue
+
+            root_input = other_inputs[0]
+            children = input_name_to_nodes[root_input]
+            children_types = [child.op_type for child in children]
+            if children_types.count('MatMul') != 3:
                 continue
 
             (add_qkv, matmul_qkv, reshape_qkv, transpose_qkv, matmul_qkv) = qkv_nodes
@@ -988,10 +990,17 @@ class BertOnnxModel(OnnxModel):
                                                   output_name_to_node):
                     continue
 
-                nodes_to_remove.extend(subgraph_nodes)
-
                 weight_input = mul_node.input[1 - self.input_index(div_node.output[0], mul_node)]
                 bias_input = last_add_node.input[1 - self.input_index(mul_node.output[0], last_add_node)]
+
+                if not self.is_constant_with_specified_dimension(weight_input, 1, "layernorm weight"):
+                    continue
+
+                if not self.is_constant_with_specified_dimension(bias_input, 1, "layernorm bias"):
+                    continue
+
+                nodes_to_remove.extend(subgraph_nodes)
+
                 normalize_node = onnx.helper.make_node('LayerNormalization',
                                                        inputs=[node.input[0], weight_input, bias_input],
                                                        outputs=[last_add_node.output[0]])
@@ -1051,7 +1060,9 @@ class BertOnnxModel(OnnxModel):
         self.fuse_reshape()
 
         self.fuse_skip_layer_norm()
+
         self.fuse_attention()
+
         self.fuse_embed_layer()
 
         # Fuse Gelu and Add Bias before it.

onnxruntime/python/tools/bert/BertOnnxModelTF.py

@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
 class BertOnnxModelTF(BertOnnxModel):
 
     def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
-        super().__init__(model, num_heads, hidden_size, sequence_length)
+        super().__init__(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
 
     """
      Fuse Gelu with Erf into one node:

onnxruntime/python/tools/bert/Gpt2OnnxModel.py

@@ -0,0 +1,12 @@
+#-------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+#--------------------------------------------------------------------------
+
+from BertOnnxModel import BertOnnxModel
+
+
+class Gpt2OnnxModel(BertOnnxModel):
+
+    def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
+        super().__init__(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)

onnxruntime/python/tools/bert/OnnxModel.py

@@ -233,7 +233,8 @@ class OnnxModel:
                                                exclude=[],
                                                return_indice=return_indice)
             if matched_parent is None:
-                logger.debug(f"Failed to match index={i} parent_input_index={parent_input_index[i]} op_type={op_type}")
+                logger.debug(f"Failed to match index={i} parent_input_index={parent_input_index[i]} op_type={op_type}",
+                             stack_info=True)
                 return None
 
             matched_parents.append(matched_parent)
@@ -304,6 +305,18 @@ class OnnxModel:
 
         return -1
 
+    def is_constant_with_specified_dimension(self, output_name, dimensions, description):
+        value = self.get_constant_value(output_name)
+        if value is None:
+            logger.debug(f"{description} {output_name} is not initializer.")
+            return False
+
+        if len(value.shape) != dimensions:
+            logger.debug(f"{description} {output_name} shall have {dimensions} dimensions. Got shape {value.shape}")
+            return False
+
+        return True
+
     def has_constant_input(self, node, expected_value, delta=0.000001):
         return self.find_constant_input(node, expected_value, delta) >= 0
 

onnxruntime/python/tools/bert/benchmark_gpt2.py

@@ -201,7 +201,7 @@ def main():
     if args.enable_optimization:
         from bert_model_optimization import optimize_model
         m = optimize_model(export_model_path,
-                           model_type='bert',
+                           model_type='gpt2',
                            gpu_only=False,
                            num_heads=12,
                            hidden_size=768,
@@ -236,5 +236,6 @@ def main():
             logger.info('PyTorch and OnnxRuntime layer {} state (present_{}) are close:'.format(layer, layer),
                         numpy.allclose(ort_outputs[1 + layer], outputs[1][layer].cpu(), rtol=1e-05, atol=1e-04))
 
+
 if __name__ == '__main__':
     main()

onnxruntime/python/tools/bert/bert_model_optimization.py

@@ -23,6 +23,7 @@
 #      PyTorch 1.2 or above, and exported to Onnx using opset version 10 or 11.
 
 import logging
+import coloredlogs
 import onnx
 import os
 import sys
@@ -33,6 +34,7 @@ from onnx import ModelProto, TensorProto, numpy_helper
 from BertOnnxModel import BertOnnxModel
 from BertOnnxModelTF import BertOnnxModelTF
 from BertOnnxModelKeras import BertOnnxModelKeras
+from Gpt2OnnxModel import Gpt2OnnxModel
 
 logger = logging.getLogger('')
 
@@ -40,7 +42,8 @@ logger = logging.getLogger('')
 MODEL_CLASSES = {
     "bert": (BertOnnxModel, "pytorch", True),
     "bert_tf": (BertOnnxModelTF, "tf2onnx", False),
-    "bert_keras": (BertOnnxModelKeras, "keras2onnx", False)
+    "bert_keras": (BertOnnxModelKeras, "keras2onnx", False),
+    "gpt2": (Gpt2OnnxModel, "pytorch", True)
 }
 
 
@@ -180,24 +183,17 @@ def optimize_model(input,
     return bert_model
 
 
+def setup_logger(verbose):
+    if verbose:
+        coloredlogs.install(level='DEBUG', fmt='[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s')
+    else:
+        coloredlogs.install(fmt='%(funcName)20s: %(message)s')
+
+
 def main():
     args = parse_arguments()
 
-    # output logging to stdout
-    log_handler = logging.StreamHandler(sys.stdout)
-    if args.verbose:
-        log_handler.setFormatter(logging.Formatter('[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s'))
-        logging_level = logging.DEBUG
-    else:
-        log_handler.setFormatter(logging.Formatter('%(filename)20s: %(message)s'))
-        logging_level = logging.INFO
-    log_handler.setLevel(logging_level)
-
-    # Avoid duplicated handlers when runing this script in multiple cells of Jupyter Notebook.
-    if not logger.hasHandlers():
-        logger.addHandler(log_handler)
-
-    logger.setLevel(logging_level)
+    setup_logger(args.verbose)
 
     bert_model = optimize_model(args.input, args.model_type, args.gpu_only, args.num_heads, args.hidden_size,
                                 args.sequence_length, args.input_int32, args.float16, args.opt_level)

onnxruntime/python/tools/bert/bert_perf_test.py

@@ -44,13 +44,26 @@ def create_session(model_path, use_gpu, intra_op_num_threads, graph_optimization
     else:
         execution_providers = ['CPUExecutionProvider'
                               ] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider']
+
         sess_options = onnxruntime.SessionOptions()
         sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
+
         if graph_optimization_level is None:
             sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        elif graph_optimization_level == 0:
+            sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+        elif graph_optimization_level == 1:
+            sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
+        elif graph_optimization_level == 2:
+            sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+        elif graph_optimization_level == 99:
+            sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
         else:
             sess_options.graph_optimization_level = graph_optimization_level
-        sess_options.intra_op_num_threads = intra_op_num_threads
+
+        if intra_op_num_threads is not None:
+            sess_options.intra_op_num_threads = intra_op_num_threads
+
         session = onnxruntime.InferenceSession(model_path, sess_options, providers=execution_providers)
 
     if use_gpu:
@@ -120,14 +133,15 @@ def setup_openmp_environ(omp_num_threads, omp_wait_policy):
 
 
 def run_one_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                 contiguous, intra_op_num_threads, omp_num_threads, omp_wait_policy, no_warmup, extra_latency):
+                 contiguous, intra_op_num_threads, omp_num_threads, omp_wait_policy, no_warmup, opt_level,
+                 extra_latency):
     # Environment variable shall be set before import onnxruntime.
     setup_openmp_environ(omp_num_threads, omp_wait_policy)
 
     test_setting = "batch_size={},sequence_length={},test_cases={},test_times={},contiguous={},use_gpu={},warmup={}".format(
         batch_size, sequence_length, test_cases, test_times, contiguous, use_gpu, not no_warmup)
 
-    session = create_session(model_path, use_gpu, intra_op_num_threads)
+    session = create_session(model_path, use_gpu, intra_op_num_threads, opt_level)
     output_names = [output.name for output in session.get_outputs()]
 
     key = to_string(model_path, session, test_setting)
@@ -160,38 +174,39 @@ def run_one_test(perf_results, model_path, all_inputs, batch_size, sequence_leng
 
 
 def launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                contiguous, intra_op_num_threads, omp_num_threads, omp_wait_policy, no_warmup, extra_latency):
+                contiguous, intra_op_num_threads, omp_num_threads, omp_wait_policy, no_warmup, opt_level,
+                extra_latency):
     process = multiprocessing.Process(target=run_one_test,
                                       args=(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu,
                                             test_cases, test_times, contiguous, intra_op_num_threads, omp_num_threads,
-                                            omp_wait_policy, no_warmup, extra_latency))
+                                            omp_wait_policy, no_warmup, opt_level, extra_latency))
     process.start()
     process.join()
 
 
 def run_perf_tests(perf_results, model_path, batch_size, sequence_length, use_gpu, test_cases, test_times, contiguous,
-                   all_inputs, test_all, no_warmup, extra_latency):
+                   all_inputs, test_all, no_warmup, opt_level, extra_latency):
     cpu_count = psutil.cpu_count(logical=False)
     logical_cores = psutil.cpu_count(logical=True)
 
     # Test a setting without any setting as baseline 1.
     launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                contiguous, None, None, None, no_warmup, extra_latency)
+                contiguous, None, None, None, no_warmup, opt_level, extra_latency)
 
     if not use_gpu:
         # For CPU: intra_op_num_threads = 1, omp_num_threads=None, omp_wait_policy=None
         # Another setting without environment variable as baseline 2.
         launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                    contiguous, 1, None, None, no_warmup, extra_latency)
+                    contiguous, 1, None, None, no_warmup, opt_level, extra_latency)
     else:
         # For GPU, we test two more settings by default:
         # (1) intra_op_num_threads = 1, omp_num_threads=cpu_count, omp_wait_policy=PASSIVE
         # (2) intra_op_num_threads = logical_cores, omp_num_threads=1, omp_wait_policy=ACTIVE
         launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                    contiguous, 1, cpu_count, 'PASSIVE', no_warmup, extra_latency)
+                    contiguous, 1, cpu_count, 'PASSIVE', no_warmup, opt_level, extra_latency)
 
         launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases, test_times,
-                    contiguous, logical_cores, 1, 'ACTIVE', no_warmup, extra_latency)
+                    contiguous, logical_cores, 1, 'ACTIVE', no_warmup, opt_level, extra_latency)
 
     # GPU latency is not sensitive to these settings. No need to test many combinations.
     # Skip remaining settings for GPU without --all flag.
@@ -221,11 +236,11 @@ def run_perf_tests(perf_results, model_path, batch_size, sequence_length, use_gp
             for omp_wait_policy in ['ACTIVE', 'PASSIVE']:
                 launch_test(perf_results, model_path, all_inputs, batch_size, sequence_length, use_gpu, test_cases,
                             test_times, contiguous, intra_op_num_threads, omp_num_threads, omp_wait_policy, no_warmup,
-                            extra_latency)
+                            opt_level, extra_latency)
 
 
 def run_performance(perf_results, model_path, batch_size, sequence_length, use_gpu, test_cases, test_times, seed,
-                    verbose, inclusive, test_all, no_warmup):
+                    verbose, inclusive, test_all, no_warmup, opt_level):
     # Try deduce input names from model.
     input_ids, segment_ids, input_mask = get_bert_inputs(model_path)
 
@@ -253,6 +268,7 @@ def run_performance(perf_results, model_path, batch_size, sequence_length, use_g
                    all_inputs,
                    test_all,
                    no_warmup,
+                   opt_level,
                    extra_latency=0)
 
     # only test contiguous array when the --all flag is set.
@@ -275,6 +291,7 @@ def run_performance(perf_results, model_path, batch_size, sequence_length, use_g
                    all_inputs,
                    test_all,
                    no_warmup,
+                   opt_level,
                    extra_latency=contiguous_latency if inclusive else 0)
 
 
@@ -298,6 +315,14 @@ def parse_arguments():
                         default=0,
                         help="number of times to run per sample. By default, the value is 1000 / samples")
 
+    parser.add_argument(
+        '--opt_level',
+        required=False,
+        type=int,
+        choices=[0, 1, 2, 99],
+        default=99,
+        help="onnxruntime optimization level: 0 - disable all, 1 - basic, 2 - extended, 99 - enable all.")
+
     parser.add_argument('--seed',
                         required=False,
                         type=int,
@@ -341,7 +366,8 @@ def main():
 
     for batch_size in batch_size_set:
         run_performance(perf_results, args.model, batch_size, args.sequence_length, args.use_gpu, args.samples,
-                        args.test_times, args.seed, args.verbose, args.inclusive, args.all, args.no_warmup)
+                        args.test_times, args.seed, args.verbose, args.inclusive, args.all, args.no_warmup,
+                        args.opt_level)
 
     # Sort the results so that the first one has smallest latency.
     sorted_results = sorted(perf_results.items(), reverse=False, key=lambda x: x[1])