diff --git a/onnxruntime/python/tools/bert/BertOnnxModel.py b/onnxruntime/python/tools/bert/BertOnnxModel.py index 428d74c722..c04f5e5f0c 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModel.py +++ b/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. diff --git a/onnxruntime/python/tools/bert/BertOnnxModelTF.py b/onnxruntime/python/tools/bert/BertOnnxModelTF.py index 7bbbbc2574..5a41a1ea5c 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModelTF.py +++ b/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: diff --git a/onnxruntime/python/tools/bert/Gpt2OnnxModel.py b/onnxruntime/python/tools/bert/Gpt2OnnxModel.py new file mode 100644 index 0000000000..949ee27357 --- /dev/null +++ b/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) diff --git a/onnxruntime/python/tools/bert/OnnxModel.py b/onnxruntime/python/tools/bert/OnnxModel.py index 7d53847fcf..1d8e36d094 100644 --- a/onnxruntime/python/tools/bert/OnnxModel.py +++ b/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 diff --git a/onnxruntime/python/tools/bert/benchmark_gpt2.py b/onnxruntime/python/tools/bert/benchmark_gpt2.py index 75d58117bf..fe3a46791e 100644 --- a/onnxruntime/python/tools/bert/benchmark_gpt2.py +++ b/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() diff --git a/onnxruntime/python/tools/bert/bert_model_optimization.py b/onnxruntime/python/tools/bert/bert_model_optimization.py index f4f471c630..feee35b721 100644 --- a/onnxruntime/python/tools/bert/bert_model_optimization.py +++ b/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) diff --git a/onnxruntime/python/tools/bert/bert_perf_test.py b/onnxruntime/python/tools/bert/bert_perf_test.py index 040afcc0b0..2495440463 100644 --- a/onnxruntime/python/tools/bert/bert_perf_test.py +++ b/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])