From 41f1280fc91adb7676891e377b86b20ece7a54b1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 14 Jul 2021 16:00:17 -0700 Subject: [PATCH] Fix transformer optimizer (#8392) * fix a few issues --- .../python/tools/transformers/benchmark.py | 6 ++-- .../tools/transformers/benchmark_helper.py | 8 +++-- .../tools/transformers/bert_perf_test.py | 32 ++++++++++--------- .../tools/transformers/fusion_attention.py | 4 +-- .../tools/transformers/fusion_layernorm.py | 4 +-- .../transformers/gpt2_beamsearch_helper.py | 6 ++-- .../python/tools/transformers/gpt2_helper.py | 7 ++-- .../tools/transformers/huggingface_models.py | 7 ++-- .../tools/transformers/onnx_exporter.py | 6 ++-- .../python/tools/transformers/onnx_model.py | 7 +++- 10 files changed, 50 insertions(+), 37 deletions(-) diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py index 3cccb40c7e..6e5d5b98ef 100644 --- a/onnxruntime/python/tools/transformers/benchmark.py +++ b/onnxruntime/python/tools/transformers/benchmark.py @@ -127,8 +127,8 @@ def run_onnxruntime(use_gpu, model_names, model_class, precision, num_threads, b continue 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, - config, input_value_type) + ort_inputs = create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names, config, + input_value_type) result_template = { "engine": "onnxruntime", "version": onnxruntime.__version__, @@ -333,7 +333,7 @@ def run_tensorflow(use_gpu, model_names, model_class, precision, num_threads, ba @run_with_tf_optimizations(do_eager_mode=False, use_xla=False) def encoder_decoder_forward(): return model(input_ids, decoder_input_ids=input_ids, training=False) - + @run_with_tf_optimizations(do_eager_mode=False, use_xla=False) def lxmert_forward(): feats = tf.random.normal([1, 1, config.visual_feat_dim]) diff --git a/onnxruntime/python/tools/transformers/benchmark_helper.py b/onnxruntime/python/tools/transformers/benchmark_helper.py index 46a5f0dd90..3d43a309f8 100644 --- a/onnxruntime/python/tools/transformers/benchmark_helper.py +++ b/onnxruntime/python/tools/transformers/benchmark_helper.py @@ -30,10 +30,13 @@ class Precision(Enum): def __str__(self): return self.value + IO_BINDING_DATA_TYPE_MAP = { "float32": numpy.float32, - # TODO: Add more. + # TODO: Add more. } + + def create_onnxruntime_session(onnx_model_path, use_gpu, enable_all_optimization=True, @@ -217,7 +220,8 @@ def inference_ort_with_io_binding(ort_session, # Bind inputs to device for name in ort_inputs.keys(): np_input = torch.from_numpy(ort_inputs[name]).to(device) - input_type = IO_BINDING_DATA_TYPE_MAP[str(ort_inputs[name].dtype)] if str(ort_inputs[name].dtype) in IO_BINDING_DATA_TYPE_MAP else data_type + input_type = IO_BINDING_DATA_TYPE_MAP[str(ort_inputs[name].dtype)] if str( + ort_inputs[name].dtype) in IO_BINDING_DATA_TYPE_MAP else data_type io_binding.bind_input(name, np_input.device.type, 0, input_type, np_input.shape, np_input.data_ptr()) # Bind outputs buffers with the sizes needed if not allocated already if len(output_buffers) == 0: diff --git a/onnxruntime/python/tools/transformers/bert_perf_test.py b/onnxruntime/python/tools/transformers/bert_perf_test.py index ff7738d4a8..45682eb18a 100644 --- a/onnxruntime/python/tools/transformers/bert_perf_test.py +++ b/onnxruntime/python/tools/transformers/bert_perf_test.py @@ -40,6 +40,7 @@ class TestSetting: seed: int verbose: bool + @dataclass class ModelSetting: model_path: str @@ -60,7 +61,8 @@ def create_session(model_path, use_gpu, intra_op_num_threads, graph_optimization if intra_op_num_threads is None and graph_optimization_level is None: session = onnxruntime.InferenceSession(model_path) else: - execution_providers = ['CPUExecutionProvider'] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider'] + execution_providers = ['CPUExecutionProvider' + ] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider'] sess_options = onnxruntime.SessionOptions() sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL @@ -103,6 +105,7 @@ def onnxruntime_inference(session, all_inputs, output_names): latency_list.append(latency) return results, latency_list + def to_string(model_path, session, test_setting): sess_options = session.get_session_options() option = "model={},".format(os.path.basename(model_path)) @@ -112,6 +115,7 @@ def to_string(model_path, session, test_setting): option += f"batch_size={test_setting.batch_size},sequence_length={test_setting.sequence_length},test_cases={test_setting.test_cases},test_times={test_setting.test_times},use_gpu={test_setting.use_gpu}" return option + def run_one_test(model_setting, test_setting, perf_results, all_inputs, intra_op_num_threads): session = create_session(model_setting.model_path, test_setting.use_gpu, intra_op_num_threads, model_setting.opt_level) @@ -148,7 +152,8 @@ def run_one_test(model_setting, test_setting, perf_results, all_inputs, intra_op def launch_test(model_setting, test_setting, perf_results, all_inputs, intra_op_num_threads): process = multiprocessing.Process(target=run_one_test, - args=(model_setting, test_setting, perf_results, all_inputs, intra_op_num_threads)) + args=(model_setting, test_setting, perf_results, all_inputs, + intra_op_num_threads)) process.start() process.join() @@ -169,7 +174,8 @@ def run_perf_tests(model_setting, test_setting, perf_results, all_inputs): for intra_op_num_threads in candidate_threads: launch_test(model_setting, test_setting, perf_results, all_inputs, intra_op_num_threads) - + + def run_performance(model_setting, test_setting, perf_results): input_ids, segment_ids, input_mask = get_bert_inputs(model_setting.model_path, model_setting.input_ids_name, model_setting.segment_ids_name, model_setting.input_mask_name) @@ -195,7 +201,8 @@ def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--model', required=True, type=str, help="bert onnx model path") - parser.add_argument('-b', '--batch_size', + parser.add_argument('-b', + '--batch_size', required=True, type=int, nargs="+", @@ -205,7 +212,8 @@ def parse_arguments(): parser.add_argument('--samples', required=False, type=int, default=10, help="number of samples to be generated") - parser.add_argument('-t', '--test_times', + parser.add_argument('-t', + '--test_times', required=False, type=int, default=0, @@ -231,7 +239,8 @@ def parse_arguments(): parser.add_argument('--use_gpu', required=False, action='store_true', help="use GPU") parser.set_defaults(use_gpu=False) - parser.add_argument('-n', '--intra_op_num_threads', + parser.add_argument('-n', + '--intra_op_num_threads', required=False, type=int, default=None, @@ -266,15 +275,8 @@ def main(): args.opt_level) for batch_size in batch_size_set: - test_setting = TestSetting( - batch_size, - args.sequence_length, - args.samples, - args.test_times, - args.use_gpu, - args.intra_op_num_threads, - args.seed, - args.verbose) + test_setting = TestSetting(batch_size, args.sequence_length, args.samples, args.test_times, args.use_gpu, + args.intra_op_num_threads, args.seed, args.verbose) print("test setting", test_setting) run_performance(model_setting, test_setting, perf_results) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 45bd8c49e9..c18225d6e9 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -119,10 +119,10 @@ class FusionAttention(Fusion): hidden_size = num_heads * head_size if self.num_heads > 0 and num_heads != self.num_heads: - logger.warn("--num_heads is {self.num_heads}. Detected value is {num_heads}. Using detected value.") + logger.warn(f"--num_heads is {self.num_heads}. Detected value is {num_heads}. Using detected value.") if self.hidden_size > 0 and hidden_size != self.hidden_size: - logger.warn("--hidden_size is {self.hidden_size}. Detected value is {hidden_size}. Using detected value.") + logger.warn(f"--hidden_size is {self.hidden_size}. Detected value is {hidden_size}. Using detected value.") return num_heads, hidden_size diff --git a/onnxruntime/python/tools/transformers/fusion_layernorm.py b/onnxruntime/python/tools/transformers/fusion_layernorm.py index 57c110dd64..b3e2f608fb 100644 --- a/onnxruntime/python/tools/transformers/fusion_layernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_layernorm.py @@ -220,7 +220,7 @@ class FusionLayerNormalizationTF(Fusion): fused_node = helper.make_node('LayerNormalization', inputs=[mul_node_3.input[0], weight_input, bias_input], outputs=[node.output[0]], - name=self.model.create_node_name("LayerNormalization", - name_prefix="SkipLayerNorm")) + name=self.model.create_node_name("LayerNormalization", name_prefix="LayerNorm")) fused_node.attribute.extend([helper.make_attribute("epsilon", float(epsilon))]) self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name diff --git a/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py b/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py index 0001785a1c..2570673692 100644 --- a/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py +++ b/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py @@ -612,13 +612,13 @@ class Gpt2BeamSearchHelper(Gpt2Helper): # add dynamic output axes present_axes = {1: 'batch_size', 3: 'cur_seq_len'} - dynamic_axes["last_state"] = {0: 'batch_size', 1: 'beam_size'} + dynamic_axes["last_state"] = {0: 'batch_size', 1: 'beam_size'} for i in range(num_layer): dynamic_axes["present_" + str(i)] = present_axes dynamic_axes["output_selected_indices"] = {0: "batch_size", 1: "'beam_size_or_1'"} - dynamic_axes["output_log_probs"] = {0: "batch_size", 1: "'beam_size'"} - dynamic_axes["output_unfinished_sents"] = {0: "batch_size", 1: "'beam_size'"} + dynamic_axes["output_log_probs"] = {0: "batch_size", 1: "'beam_size'"} + dynamic_axes["output_unfinished_sents"] = {0: "batch_size", 1: "'beam_size'"} dynamic_axes["current_step_results"] = {0: "beam_size_or_1", 1: "total_seq_len"} dynamic_axes["current_step_scores"] = {0: "beam_size_or_1", 1: "total_seq_len"} diff --git a/onnxruntime/python/tools/transformers/gpt2_helper.py b/onnxruntime/python/tools/transformers/gpt2_helper.py index 83142d5e54..73f37cc415 100644 --- a/onnxruntime/python/tools/transformers/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/gpt2_helper.py @@ -33,6 +33,7 @@ class GPT2ModelNoPastState(GPT2Model): def forward(self, input_ids): return super().forward(input_ids, use_cache=False, return_dict=False) + class TFGPT2ModelNoPastState(TFGPT2Model): """ Here we wrap a class to disable past state output. """ @@ -43,6 +44,7 @@ class TFGPT2ModelNoPastState(TFGPT2Model): def forward(self, input_ids): return super().call(input_ids, use_cache=False) + class MyGPT2Model(GPT2Model): """ Here we wrap a class for Onnx model conversion for GPT2Model with past state. """ @@ -52,7 +54,8 @@ class MyGPT2Model(GPT2Model): @staticmethod def post_process(result, num_layer): if isinstance(result[1][0], tuple) or isinstance(result[1][0], list): - assert len(result[1]) == num_layer and len(result[1][0]) == 2 #and len(result[1][0][0].shape) == 4 and result[1][0][0].shape == result[1][0][1].shape + assert len(result[1]) == num_layer and len(result[1][0]) == 2 + #assert len(result[1][0][0].shape) == 4 and result[1][0][0].shape == result[1][0][1].shape present = [] for i in range(num_layer): # Since transformers v4.*, past key and values are separated outputs. @@ -321,7 +324,7 @@ class Gpt2Helper: logger.info( f"Shapes: input_ids={dummy_inputs.input_ids.shape} past={dummy_inputs.past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}" ) - + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) torch.onnx.export(model, diff --git a/onnxruntime/python/tools/transformers/huggingface_models.py b/onnxruntime/python/tools/transformers/huggingface_models.py index a81cdb82da..a2cd823e9f 100644 --- a/onnxruntime/python/tools/transformers/huggingface_models.py +++ b/onnxruntime/python/tools/transformers/huggingface_models.py @@ -39,9 +39,8 @@ MODELS = { "gpt2-large": (["input_ids"], 11, True, "gpt2"), "gpt2-xl": (["input_ids"], 11, True, "gpt2"), "distilgpt2": (["input_ids"], 11, False, "gpt2"), - # Transformer-XL - "transfo-xl-wt103": - (["input_ids", "mems"], 12, False, "bert"), # Models uses Einsum, which need opset version 12 and PyTorch 1.5.0 or above. + # Transformer-XL (Models uses Einsum, which need opset version 12 or later.) + "transfo-xl-wt103": (["input_ids", "mems"], 12, False, "bert"), # XLNet "xlnet-base-cased": (["input_ids"], 12, False, "bert"), "xlnet-large-cased": (["input_ids"], 12, False, "bert"), @@ -121,7 +120,7 @@ MODELS = { # "funnel-transformer/large": (["input_ids"], 12, True, "bert"), # "funnel-transformer/large-base": (["input_ids"], 12, True, "bert"), # "funnel-transformer/xlarge": (["input_ids"], 12, True, "bert"), - # "funnel-transformer/xlarge-base": (["input_ids"], 12, True, "bert"), + # "funnel-transformer/xlarge-base": (["input_ids"], 12, True, "bert"), # Layoutlm "microsoft/layoutlm-base-uncased": (["input_ids"], 11, False, "bert"), "microsoft/layoutlm-large-uncased": (["input_ids"], 11, False, "bert"), diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 1aaa5144f2..d46586f0e7 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -436,13 +436,13 @@ def export_onnx_model_from_tf(model_name, opset_version, use_external_data_forma example_outputs = model(example_inputs, training=False) output_names = None - - # For xlnet models, only compare the last_hidden_state output. + + # For xlnet models, only compare the last_hidden_state output. if model_name == "xlnet-base-cased" or model_name == "xlnet-large-cased": output_names = ["last_hidden_state"] example_outputs = example_outputs["last_hidden_state"] - # Flatten is needed for gpt2 and distilgpt2. Output name sorting is needed for tf2onnx outputs to match onnx outputs. + # Flatten is needed for gpt2 and distilgpt2. Output name sorting is needed for tf2onnx outputs to match onnx outputs. from tensorflow.python.util import nest example_outputs_flatten = nest.flatten(example_outputs) diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index be74dd9218..d0292ba7db 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -495,7 +495,12 @@ class OnnxModel: from packaging.version import Version import onnxconverter_common as oc if Version(oc.__version__) > Version("1.7.0"): - self.model = oc.float16.convert_float_to_float16(self.model, keep_io_types=cast_input_output) + # Use symbolic shape inference since custom operators (like Gelu, SkipLayerNormalization etc) are not recognized by onnx shape inference. + shape_infer_helper = SymbolicShapeInferenceHelper(self.model) + model_with_shape = shape_infer_helper.infer_shapes(self.model, auto_merge=True, guess_output_rank=False) + self.model = oc.float16.convert_float_to_float16(model_with_shape, + keep_io_types=cast_input_output, + disable_shape_infer=True) return graph = self.model.graph