diff --git a/onnxruntime/python/tools/bert/BertOnnxModel.py b/onnxruntime/python/tools/bert/BertOnnxModel.py index a23cb0da8e..3917730a7e 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModel.py +++ b/onnxruntime/python/tools/bert/BertOnnxModel.py @@ -15,6 +15,19 @@ from OnnxModel import OnnxModel logger = logging.getLogger(__name__) +class BertOptimizationOptions: + + def __init__(self, model_type): + self.enable_attention = True + self.enable_skip_layer_norm = True + self.enable_embed_layer_norm = True + self.enable_bias_skip_layer_norm = True + self.enable_bias_gelu = True + + if model_type == 'gpt2': + self.enable_skip_layer_norm = False + + class BertOnnxModel(OnnxModel): def __init__(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only): @@ -38,15 +51,31 @@ class BertOnnxModel(OnnxModel): self.bert_inputs = [] + def cast_input_to_int32(self, input_name): + cast_output = input_name + '_int32' + + # Avoid consequent Cast nodes. + inputs = [input_name] + output_name_to_node = self.output_name_to_node() + if input_name in output_name_to_node: + parent_node = output_name_to_node[input_name] + if parent_node and parent_node.op_type == 'Cast': + inputs = [parent_node.input[0]] + + cast_node = onnx.helper.make_node('Cast', inputs=inputs, outputs=[cast_output]) + cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.INT32))]) + self.add_node(cast_node) + + return cast_output, cast_node + 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: - cast_output = input_name + '_int32' - cast_node = onnx.helper.make_node('Cast', inputs=[input_name], outputs=[cast_output]) - cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.INT32))]) - self.add_node(cast_node) + cast_output, cast_node = self.cast_input_to_int32(input_name) + logger.debug("Casted graph input {input_name} to int32") return True, cast_output + logger.debug(f"Did not cast graph input {input_name} to int32: found {graph_input is not None}") return False, input_name def undo_cast_input_to_int32(self, input_name): @@ -69,7 +98,12 @@ class BertOnnxModel(OnnxModel): return self.mask_indice[input] # Add cast to convert int64 to int32 - casted, input_name = self.cast_graph_input_to_int32(input) + if self.find_graph_input(input): + casted, input_name = self.cast_graph_input_to_int32(input) + else: + input_name, cast_node = self.cast_input_to_int32(input) + casted = True + if casted: self.mask_casted[input] = input_name @@ -125,19 +159,12 @@ class BertOnnxModel(OnnxModel): vals=qkv_weight.flatten().tolist()) self.add_initializer(weight) - weight_input = onnx.helper.make_tensor_value_info(weight.name, TensorProto.FLOAT, - [self.hidden_size, 3 * self.hidden_size]) - self.add_input(weight_input) - bias = onnx.helper.make_tensor(name=attention_node_name + '_qkv_bias', data_type=TensorProto.FLOAT, dims=[3 * self.hidden_size], vals=qkv_bias.flatten().tolist()) self.add_initializer(bias) - bias_input = onnx.helper.make_tensor_value_info(bias.name, TensorProto.FLOAT, [3 * self.hidden_size]) - self.add_input(bias_input) - attention_node = onnx.helper.make_node( 'Attention', inputs=[input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias', mask_index], @@ -350,6 +377,7 @@ class BertOnnxModel(OnnxModel): """ def fuse_gelu_with_tanh(self): + logger.debug(f"start FastGelu fusion...") input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() @@ -435,19 +463,23 @@ class BertOnnxModel(OnnxModel): nodes_to_add.append(gelu_node) if len(nodes_to_add) > 0: - logger.info("Fused FastGelu count: {len(nodes_to_add)}") + logger.info(f"Fused FastGelu count: {len(nodes_to_add)}") self.remove_nodes(nodes_to_remove) self.add_nodes(nodes_to_add) - def fuse_add_bias_gelu(self): + def fuse_bias_gelu(self, is_fastgelu): + gelu_op_type = 'FastGelu' if is_fastgelu else 'Gelu' + bias_gelu_op_type = 'FastGelu' if is_fastgelu else 'BiasGelu' + logger.debug(f"start Bias and {gelu_op_type} fusion...") + input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() nodes_to_remove = [] nodes_to_add = [] # Don't need to fuse Gelu+Add here because ORT native code can handle it - for node in self.get_nodes_by_op_type('FastGelu'): + for node in self.get_nodes_by_op_type(gelu_op_type): if len(node.input) != 1: continue @@ -476,20 +508,21 @@ class BertOnnxModel(OnnxModel): continue nodes_to_remove.extend(subgraph_nodes) - gelu_node = onnx.helper.make_node('FastGelu', + gelu_node = onnx.helper.make_node(bias_gelu_op_type, inputs=[matmul.output[0], add.input[bias_index]], outputs=node.output, - name=self.create_node_name('FastGelu', "FastGelu_AddBias_")) + name=self.create_node_name(bias_gelu_op_type, gelu_op_type + "_AddBias_")) gelu_node.domain = "com.microsoft" nodes_to_add.append(gelu_node) if len(nodes_to_add) > 0: - logger.info(f"Fused FastGelu with Bias count:{len(nodes_to_add)}") + logger.info(f"Fused {bias_gelu_op_type} with Bias count:{len(nodes_to_add)}") self.remove_nodes(nodes_to_remove) self.add_nodes(nodes_to_add) def fuse_add_bias_skip_layer_norm(self): + logger.debug(f"start Bias and SkipLayerNormalization fusion...") input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() nodes_to_remove = [] @@ -552,6 +585,7 @@ class BertOnnxModel(OnnxModel): self.add_nodes(nodes_to_add) def fuse_reshape(self): + logger.debug(f"start Reshape fusion...") nodes = self.nodes() input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() @@ -705,21 +739,12 @@ class BertOnnxModel(OnnxModel): SkipLayerNormalization """ - def fuse_embed_layer(self): + def fuse_embed_layer_without_mask(self): + logger.debug(f"start EmbedLayerNormalization (no mask) fusion...") nodes = self.nodes() input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() - if len(self.mask_indice) == 0: - logger.info("skip embed layer fusion since mask input is not found") - return - if len(self.mask_indice) > 1: - logger.info("skip embed layer fusion since there are multiple mask inputs found") - return - mask_input_name = next(iter(self.mask_indice)) - mask_output_name = self.mask_indice[mask_input_name] - mask_node = output_name_to_node[mask_output_name] - nodes_to_remove = [] # Find the first normalize node could be embedding layer. @@ -731,6 +756,14 @@ class BertOnnxModel(OnnxModel): if self.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) is not None: normalize_node = node break + # In case user disables attention fusion, check whether subgraph looks like Attention. + if node.output[0] not in input_name_to_nodes: + continue + children = input_name_to_nodes[node.output[0]] + children_types = sorted([child.op_type for child in children]) + if children_types == ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']: + normalize_node = node + break if normalize_node is None: if len(self.get_nodes_by_op_type("EmbedLayerNormalization")) == 0: @@ -780,7 +813,9 @@ class BertOnnxModel(OnnxModel): segment_ids = segment_embedding_gather.input[1] if position_embedding_expand: - subgraph_nodes = self.get_parent_subgraph_nodes(position_embedding_expand, [], output_name_to_node) + input_parent = self.get_parent(position_embedding_shape, 0, output_name_to_node) + subgraph_nodes = self.get_parent_subgraph_nodes(position_embedding_expand, + [input_parent] if input_parent else [], output_name_to_node) nodes_to_remove.extend(subgraph_nodes) nodes_to_remove.extend(word_embedding_path) @@ -788,22 +823,37 @@ class BertOnnxModel(OnnxModel): nodes_to_remove.extend(segment_embedding_path) nodes_to_remove.extend([normalize_node]) - nodes_to_remove.extend([mask_node]) # store inputs for further processing - self.bert_inputs = [input_ids, segment_ids, mask_input_name] + if self.find_graph_input(input_ids): + self.bert_inputs = [input_ids, segment_ids] if self.find_graph_input(segment_ids) else [input_ids] - if not self.input_int32: - # When mask has been casted to int32, use that casted one as input of embed layer norm. - if mask_input_name in self.mask_casted: - mask_input_name = self.mask_casted[mask_input_name] - - # Cast input_ids and segment_ids to int32. - casted, input_ids = self.cast_graph_input_to_int32(input_ids) - - casted, segment_ids = self.cast_graph_input_to_int32(segment_ids) + # Cast input_ids and segment_ids to int32. + if self.find_graph_input(input_ids): + if not self.input_int32: + casted, input_ids = self.cast_graph_input_to_int32(input_ids) else: - self.undo_cast_input_to_int32(mask_input_name) + input_ids, input_ids_cast_node = self.cast_input_to_int32(input_ids) + + if self.find_graph_input(segment_ids): + if not self.input_int32: + casted, segment_ids = self.cast_graph_input_to_int32(segment_ids) + else: + segment_ids, segment_ids_cast_node = self.cast_input_to_int32(segment_ids) + + segment_id_path = self.match_parent_path( + segment_ids_cast_node, ['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], + [0, 0, 1, 0, 0, 0]) + if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]: + logger.debug("Simplify semgent id path...") + self.add_node( + onnx.helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"])) + self.add_node( + onnx.helper.make_node('ConstantOfShape', + inputs=["input_shape"], + outputs=["zeros_for_input_shape"], + value=onnx.helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1]))) + segment_ids = "zeros_for_input_shape" embed_node = onnx.helper.make_node( 'EmbedLayerNormalization', @@ -814,10 +864,9 @@ class BertOnnxModel(OnnxModel): position_embedding_weight_node.input[0], segment_embedding_gather.input[0], normalize_node.input[2], - normalize_node.input[3], # gamma and beta - mask_input_name + normalize_node.input[3] # gamma and beta ], - outputs=["embed_output", mask_output_name], + outputs=["embed_output", "dummy_mask_index"], name="EmbedLayer") embed_node.domain = "com.microsoft" @@ -826,13 +875,50 @@ class BertOnnxModel(OnnxModel): self.remove_nodes(nodes_to_remove) self.add_node(embed_node) - self.update_graph() - logger.info("Fused EmbedLayerNormalization count: 1") + self.prune_graph() + + return embed_node + + def fuse_embed_layer(self): + embed_node = self.fuse_embed_layer_without_mask() + if embed_node is None: + logger.info("Fused EmbedLayerNormalization count: 0") + return + + if len(self.mask_indice) > 1: + logger.info("There are multiple mask inputs found!") + + if len(self.mask_indice) != 1: + logger.info("Fused EmbedLayerNormalization (no mask) count: 1") + return + + mask_input_name = next(iter(self.mask_indice)) + mask_output_name = self.mask_indice[mask_input_name] + mask_node = output_name_to_node[mask_output_name] + + nodes_to_remove = [] + nodes_to_remove.extend([mask_node]) + + # store inputs for further processing + self.bert_inputs.append(mask_input_name) + + if not self.input_int32: + # When mask has been casted to int32, use that casted one as input of embed layer norm. + if mask_input_name in self.mask_casted: + mask_input_name = self.mask_casted[mask_input_name] + else: + self.undo_cast_input_to_int32(mask_input_name) + + embed_node.input[7] = mask_input_name + embed_node.output[1] = mask_output_name + logger.info("Added mask to EmbedLayerNormalization") # Change graph input data type int32 if needed. if self.input_int32: self.change_input_to_int32() + logger.info("Fused EmbedLayerNormalization count: 1") + def get_bert_inputs(self, include_mask=True): return self.bert_inputs if include_mask else self.bert_inputs[:2] @@ -856,14 +942,13 @@ class BertOnnxModel(OnnxModel): graph = self.graph() batch_size = self.get_batch_size_from_graph_input() - input_batch_size = batch_size if isinstance(batch_size, int) else 1 new_graph_inputs = [] bert_inputs = self.get_bert_inputs() for input in graph.input: if input.name in bert_inputs: - int32_input = onnx.helper.make_tensor_value_info(input.name, TensorProto.INT32, - [input_batch_size, self.sequence_length]) + input_shape = [batch_size if isinstance(batch_size, int) else 1, self.sequence_length] + int32_input = onnx.helper.make_tensor_value_info(input.name, TensorProto.INT32, input_shape) new_graph_inputs.append(int32_input) else: new_graph_inputs.append(input) @@ -924,6 +1009,8 @@ class BertOnnxModel(OnnxModel): | | +----------------------+ """ + logger.debug(f"start LayerNormalization fusion...") + input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() @@ -1013,6 +1100,8 @@ class BertOnnxModel(OnnxModel): """ Fuse Add + LayerNormalization into one node: SkipLayerNormalization """ + logger.debug(f"start SkipLayerNormaliation fusion...") + input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() @@ -1043,9 +1132,9 @@ class BertOnnxModel(OnnxModel): return def postprocess(self): - return + self.prune_graph() - def optimize(self): + def optimize(self, options: BertOptimizationOptions = None): self.fuse_layer_norm() self.fuse_gelu() @@ -1054,20 +1143,28 @@ class BertOnnxModel(OnnxModel): self.fuse_reshape() - self.fuse_skip_layer_norm() + if (options is None) or options.enable_skip_layer_norm: + self.fuse_skip_layer_norm() - self.fuse_attention() + if (options is None) or options.enable_attention: + self.fuse_attention() - self.fuse_embed_layer() - - # Fuse Gelu and Add Bias before it. - self.fuse_add_bias_gelu() - - # Fuse SkipLayerNormalization and Add Bias before it. - self.fuse_add_bias_skip_layer_norm() + if (options is None) or options.enable_embed_layer_norm: + self.fuse_embed_layer() + # Post-processing like removing extra reshape nodes. self.postprocess() + # Bias fusion is done after postprocess to avoid extra Reshape between bias and Gelu/FastGelu/SkipLayerNormalization + if (options is None) or options.enable_bias_gelu: + # Fuse Gelu and Add Bias before it. + self.fuse_bias_gelu(is_fastgelu=True) + self.fuse_bias_gelu(is_fastgelu=False) + + if (options is None) or options.enable_bias_skip_layer_norm: + # Fuse SkipLayerNormalization and Add Bias before it. + self.fuse_add_bias_skip_layer_norm() + if self.float16: self.convert_model_float32_to_float16() @@ -1090,6 +1187,7 @@ class BertOnnxModel(OnnxModel): for op in ops: nodes = self.get_nodes_by_op_type(op) op_count[op] = len(nodes) + logger.info(f"Optimized operators:{op_count}") return op_count def is_fully_optimized(self): diff --git a/onnxruntime/python/tools/bert/Gpt2OnnxModel.py b/onnxruntime/python/tools/bert/Gpt2OnnxModel.py index 949ee27357..6a0633c61b 100644 --- a/onnxruntime/python/tools/bert/Gpt2OnnxModel.py +++ b/onnxruntime/python/tools/bert/Gpt2OnnxModel.py @@ -2,11 +2,189 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. #-------------------------------------------------------------------------- - +import logging +import onnx +import sys +import argparse +import numpy as np +from collections import deque +from onnx import ModelProto, TensorProto, numpy_helper from BertOnnxModel import BertOnnxModel +logger = logging.getLogger(__name__) + 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) + + def fuse_attention(self): + """ + Fuse Attention subgraph into one Attention node. + """ + logger.debug(f"start attention fusion...") + + input_name_to_nodes = self.input_name_to_nodes() + output_name_to_node = self.output_name_to_node() + + nodes_to_remove = [] + attention_count = 0 + + for normalize_node in self.get_nodes_by_op_type("LayerNormalization"): + return_indice = [] + qkv_nodes = self.match_parent_path( + normalize_node, + ['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'], + [0, None, 0, 0, 0, 0, 0], + output_name_to_node=output_name_to_node, + return_indice=return_indice + ) # yapf: disable + if qkv_nodes is None: + continue + (add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes + + another_input = add_qkv.input[1 - return_indice[0]] + + v_nodes = self.match_parent_path( + matmul_qkv, + ['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'], + [1, 0, 0, 0, 0, 0]) # yapf: disable + if v_nodes is None: + logger.debug("fuse_attention: failed to match v path") + continue + (transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes + + layernorm_before_attention = self.get_parent(reshape_before_gemm, 0, output_name_to_node) + if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization': + logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}") + continue + + if not another_input in layernorm_before_attention.input: + logger.debug("Add and LayerNormalization shall have one same input") + continue + + qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0]) + if qk_nodes is None: + logger.debug("fuse_attention: failed to match qk path") + continue + (softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes + + q_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0]) + if q_nodes is None: + logger.debug("fuse_attention: failed to match q path") + continue + (transpose_q, reshape_q, split_q) = q_nodes + if split_v != split_q: + logger.debug("fuse_attention: skip since split_v != split_q") + continue + + k_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0]) + if k_nodes is None: + logger.debug("fuse_attention: failed to match k path") + continue + (transpose_k, reshape_k, split_k) = k_nodes + if split_v != split_k: + logger.debug("fuse_attention: skip since split_v != split_k") + continue + + mask_nodes = self.match_parent_path( + sub_qk, + ['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], + [1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable + if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") + continue + (mul_mask, sub_mask, slice_mask, slice_mask_0, unsqueeze_mask, sub_mask, squeeze_mask, slice_mask_1, + shape_mask, div_mask) = mask_nodes + + if div_qk != div_mask: + logger.debug("fuse_attention: skip since div_qk != div_mask") + continue + + self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0], + attention_count == 0) + nodes_to_remove.extend([reshape_qkv, transpose_qkv, matmul_qkv]) + nodes_to_remove.extend(qk_nodes) + nodes_to_remove.extend(q_nodes) + nodes_to_remove.extend(k_nodes) + nodes_to_remove.extend(v_nodes) + nodes_to_remove.extend(mask_nodes) + attention_count += 1 + + self.remove_nodes(nodes_to_remove) + self.prune_graph() + logger.info(f"Fused Attention count:{attention_count}") + + def create_attention_node(self, gemm, gemm_qkv, input, output, add_graph_input): + attention_node_name = self.create_node_name('Attention') + attention_node = onnx.helper.make_node('Attention', + inputs=[input, gemm.input[1], gemm.input[2]], + outputs=[attention_node_name + "_output"], + name=attention_node_name) + attention_node.domain = "com.microsoft" + attention_node.attribute.extend( + [onnx.helper.make_attribute("num_heads", self.num_heads), + onnx.helper.make_attribute("unidirectional", 1)]) + + matmul_node = onnx.helper.make_node('MatMul', + inputs=[attention_node_name + "_output", gemm_qkv.input[1]], + outputs=[attention_node_name + "_matmul_output"], + name=attention_node_name + "_matmul") + + add_node = onnx.helper.make_node('Add', + inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]], + outputs=[output], + name=attention_node_name + "_add") + + self.add_node(attention_node) + self.add_node(matmul_node) + self.add_node(add_node) + + def postprocess(self): + """ + Remove extra reshape nodes. + """ + logger.debug(f"start postprocessing...") + + input_name_to_nodes = self.input_name_to_nodes() + output_name_to_node = self.output_name_to_node() + + reshape_count = 0 + for gemm_node in self.get_nodes_by_op_type("Gemm"): + reshape_after_gemm = self.find_first_child_by_type(gemm_node, + 'Reshape', + input_name_to_nodes, + recursive=False) + + return_indice = [] + nodes = self.match_parent_path(gemm_node, ['Reshape', 'FastGelu'], [0, 0], output_name_to_node) + if nodes is None: + nodes = self.match_parent_path(gemm_node, ['Reshape', 'LayerNormalization'], [0, 0], + output_name_to_node) + if nodes is None: + continue + (reshape_before_gemm, root_node) = nodes + + matmul_node_name = self.create_node_name('MatMul', 'FullyConnect_MatMul') + matmul_node = onnx.helper.make_node('MatMul', + inputs=[matmul_node_name + "_input", gemm_node.input[1]], + outputs=[matmul_node_name + "_output"], + name=matmul_node_name) + + add_node_name = self.create_node_name('Add', 'FullyConnect_Add') + add_node = onnx.helper.make_node('Add', + inputs=[matmul_node_name + "_output", gemm_node.input[2]], + outputs=[add_node_name + "_output"], + name=add_node_name) + + root_node.output[0] = matmul_node_name + "_input" + self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output") + + self.add_node(matmul_node) + self.add_node(add_node) + + reshape_count += 2 + + self.prune_graph() + logger.info(f"Remove Reshape count:{reshape_count}") diff --git a/onnxruntime/python/tools/bert/benchmark_gpt2.py b/onnxruntime/python/tools/bert/benchmark_gpt2.py index fe3a46791e..6be7310ee0 100644 --- a/onnxruntime/python/tools/bert/benchmark_gpt2.py +++ b/onnxruntime/python/tools/bert/benchmark_gpt2.py @@ -11,10 +11,16 @@ import psutil import argparse import logging import torch -from transformers import GPT2Model, GPT2Tokenizer +from transformers import GPT2Model, GPT2LMHeadModel, GPT2Tokenizer logger = logging.getLogger('') +# Map alias to a tuple of Model Class and pretrained model name +MODEL_CLASSES = { + "gpt2": (GPT2Model, GPT2Tokenizer, "gpt2"), + "distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2") +} + def dump_environment(): if "OMP_NUM_THREADS" in os.environ: @@ -44,9 +50,7 @@ def pytorch_inference(model, input_ids, past=None, total_runs=100): with torch.no_grad(): for _ in range(total_runs): start = time.time() - outputs = model( - input_ids=input_ids, - past=past) #attention_mask=inputs['attention_mask'], token_type_ids=inputs['token_type_ids'], + outputs = model(input_ids=input_ids, past=past) latency.append(time.time() - start) logger.info("PyTorch Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f'))) @@ -88,6 +92,12 @@ def inference(model, ort_session, input_ids, past=None, total_runs=100, verify_o def parse_arguments(): parser = argparse.ArgumentParser() + parser.add_argument('--model_type', + required=True, + type=str, + choices=list(MODEL_CLASSES.keys()), + help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())) + parser.add_argument('--cache_dir', required=True, type=str, help="cache directory") parser.add_argument('--output_dir', required=True, type=str, help="output onnx model directory") @@ -132,6 +142,24 @@ def setup_logger(verbose=True): logger.setLevel(logging_level) +def remove_past_outputs(export_model_path): + from onnx import ModelProto + from OnnxModel import OnnxModel + + model = ModelProto() + with open(export_model_path, "rb") as f: + model.ParseFromString(f.read()) + bert_model = OnnxModel(model) + + # remove past state outputs and only keep the first output. + keep_output_names = [bert_model.model.graph.output[0].name] + logger.info(f"Prune graph to keep the first output and drop past state outputs:{keep_output_names}") + bert_model.prune_graph(keep_output_names) + onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_out1.onnx'.format(int(enable_past_input))) + bert_model.save_model_to_file(onnx_model_path) + return onnx_model_path + + def main(): args = parse_arguments() setup_logger(args.verbose) @@ -147,7 +175,8 @@ def main(): if not os.path.exists(output_dir): os.makedirs(output_dir) - model_class, tokenizer_class, model_name_or_path = (GPT2Model, GPT2Tokenizer, 'gpt2') + (model_class, tokenizer_class, model_name_or_path) = MODEL_CLASSES[args.model_type] + tokenizer = tokenizer_class.from_pretrained(model_name_or_path, cache_dir=cache_dir) model = model_class.from_pretrained(model_name_or_path, cache_dir=cache_dir) model.eval().cpu() @@ -197,10 +226,11 @@ def main(): setup_environment(args.use_openmp) import onnxruntime - onnx_model_path = export_model_path + onnx_model_path = export_model_path if enable_past_input else remove_past_outputs(export_model_path) + if args.enable_optimization: from bert_model_optimization import optimize_model - m = optimize_model(export_model_path, + m = optimize_model(onnx_model_path, model_type='gpt2', gpu_only=False, num_heads=12, diff --git a/onnxruntime/python/tools/bert/bert_model_optimization.py b/onnxruntime/python/tools/bert/bert_model_optimization.py index feee35b721..ec1689d531 100644 --- a/onnxruntime/python/tools/bert/bert_model_optimization.py +++ b/onnxruntime/python/tools/bert/bert_model_optimization.py @@ -31,7 +31,7 @@ import argparse import numpy as np from collections import deque from onnx import ModelProto, TensorProto, numpy_helper -from BertOnnxModel import BertOnnxModel +from BertOnnxModel import BertOnnxModel, BertOptimizationOptions from BertOnnxModelTF import BertOnnxModelTF from BertOnnxModelKeras import BertOnnxModelKeras from Gpt2OnnxModel import Gpt2OnnxModel @@ -137,6 +137,33 @@ def parse_arguments(): help="whether the target device is gpu or not") parser.set_defaults(gpu_only=False) + parser.add_argument('--disable_attention', required=False, action='store_true', help="disable Attention fusion") + parser.set_defaults(disable_attention=False) + + parser.add_argument('--disable_skip_layer_norm', + required=False, + action='store_true', + help="disable SkipLayerNormalization fusion") + parser.set_defaults(disable_skip_layer_norm=False) + + parser.add_argument('--disable_embed_layer_norm', + required=False, + action='store_true', + help="disable EmbedLayerNormalization fusion") + parser.set_defaults(disable_embed_layer_norm=False) + + parser.add_argument('--disable_bias_skip_layer_norm', + required=False, + action='store_true', + help="disable Add Bias and SkipLayerNormalization fusion") + parser.set_defaults(disable_bias_skip_layer_norm=False) + + parser.add_argument('--disable_bias_gelu', + required=False, + action='store_true', + help="disable Add Bias and Gelu/FastGelu fusion") + parser.set_defaults(disable_bias_gelu=False) + parser.add_argument('--verbose', required=False, action='store_true') parser.set_defaults(verbose=False) @@ -152,6 +179,21 @@ def parse_arguments(): return args +def get_optimization_options(args): + optimization_options = BertOptimizationOptions(args.model_type) + if args.disable_attention: + optimization_options.enable_attention = False + if args.disable_skip_layer_norm: + optimization_options.enable_skip_layer_norm = False + if args.disable_embed_layer_norm: + optimization_options.enable_embed_layer_norm = False + if args.disable_bias_skip_layer_norm: + optimization_options.enable_bias_skip_layer_norm = False + if args.disable_bias_gelu: + optimization_options.enable_bias_gelu = False + return optimization_options + + def optimize_model(input, model_type, gpu_only, @@ -160,7 +202,8 @@ def optimize_model(input, sequence_length, input_int32, float16, - opt_level=99): + opt_level=99, + optimization_options=None): (optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type] input_model_path = input @@ -177,8 +220,11 @@ def optimize_model(input, f"Model producer not matched: Expect {producer}, Got {model.producer_name} {model.producer_version}. Please specify correct --model_type parameter." ) + if optimization_options is None: + optimization_options = BertOptimizationOptions(model_type) + bert_model = optimizer_class(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only) - bert_model.optimize() + bert_model.optimize(optimization_options) return bert_model @@ -195,8 +241,11 @@ def main(): setup_logger(args.verbose) + optimization_options = get_optimization_options(args) + 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) + args.sequence_length, args.input_int32, args.float16, args.opt_level, + optimization_options) bert_model.save_model_to_file(args.output) diff --git a/onnxruntime/python/tools/bert/bert_perf_test.py b/onnxruntime/python/tools/bert/bert_perf_test.py index 2495440463..62de94635d 100644 --- a/onnxruntime/python/tools/bert/bert_perf_test.py +++ b/onnxruntime/python/tools/bert/bert_perf_test.py @@ -240,9 +240,10 @@ def run_perf_tests(perf_results, model_path, batch_size, sequence_length, use_gp def run_performance(perf_results, model_path, batch_size, sequence_length, use_gpu, test_cases, test_times, seed, - 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) + verbose, inclusive, test_all, no_warmup, opt_level, input_ids_name, segment_ids_name, + input_mask_name): + + input_ids, segment_ids, input_mask = get_bert_inputs(model_path, input_ids_name, segment_ids_name, input_mask_name) # Do not generate random mask for performance test. print(f"Generating {test_cases} samples for batch_size={batch_size} sequence_length={sequence_length}") @@ -347,6 +348,10 @@ def parse_arguments(): parser.add_argument('--no_warmup', required=False, action='store_true', help="do not use one sample for warm-up.") parser.set_defaults(no_warmup=False) + parser.add_argument('--input_ids', required=False, type=str, default=None, help="input name for input ids") + parser.add_argument('--segment_ids', required=False, type=str, default=None, help="input name for segment ids") + parser.add_argument('--input_mask', required=False, type=str, default=None, help="input name for attention mask") + args = parser.parse_args() return args @@ -367,7 +372,7 @@ 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.opt_level) + args.opt_level, args.input_ids, args.segment_ids, args.input_mask) # Sort the results so that the first one has smallest latency. sorted_results = sorted(perf_results.items(), reverse=False, key=lambda x: x[1]) diff --git a/onnxruntime/python/tools/bert/bert_test_data.py b/onnxruntime/python/tools/bert/bert_test_data.py index 5ab0f5e7da..b2c609056a 100644 --- a/onnxruntime/python/tools/bert/bert_test_data.py +++ b/onnxruntime/python/tools/bert/bert_test_data.py @@ -107,15 +107,22 @@ def fake_test_data(batch_size, sequence_length, test_cases, dictionary_size, ver """ Generate fake input data for test. """ + assert input_ids is not None + np.random.seed(random_seed) random.seed(random_seed) all_inputs = [] for test_case in range(test_cases): input_1 = fake_input_ids_data(input_ids, batch_size, sequence_length, dictionary_size) - input_2 = fake_segment_ids_data(segment_ids, batch_size, sequence_length) - input_3 = fake_input_mask_data(input_mask, batch_size, sequence_length, random_mask_length) - inputs = {input_ids.name: input_1, segment_ids.name: input_2, input_mask.name: input_3} + inputs = {input_ids.name: input_1} + + if segment_ids: + inputs[segment_ids.name] = fake_segment_ids_data(segment_ids, batch_size, sequence_length) + + if input_mask: + inputs[input_mask.name] = fake_input_mask_data(input_mask, batch_size, sequence_length, random_mask_length) + if verbose and len(all_inputs) == 0: print("Example inputs", inputs) all_inputs.append(inputs) @@ -144,7 +151,7 @@ def get_graph_input_from_embed_node(onnx_model, embed_node, input_index): return graph_input -def get_bert_inputs(onnx_file): +def get_bert_inputs(onnx_file, input_ids_name=None, segment_ids_name=None, input_mask_name=None): """ Get graph inputs for bert model. First, we will deduce from EmbedLayerNormalization node. If not found, we will guess based on naming. @@ -154,8 +161,31 @@ def get_bert_inputs(onnx_file): model.ParseFromString(f.read()) onnx_model = OnnxModel(model) - graph_inputs = onnx_model.get_graph_inputs_excluding_initializers() + + if input_ids_name is not None: + input_ids = onnx_model.find_graph_input(input_ids_name) + if input_ids is None: + raise ValueError(f"Graph does not have input named {input_ids_name}") + + segment_ids = None + if segment_ids_name: + segment_ids = onnx_model.find_graph_input(segment_ids_name) + if segment_ids is None: + raise ValueError(f"Graph does not have input named {segment_ids_name}") + + input_mask = None + if input_mask_name: + input_mask = onnx_model.find_graph_input(input_mask_name) + if input_mask is None: + raise ValueError(f"Graph does not have input named {input_mask_name}") + + expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0) + if len(graph_inputs) != expected_inputs: + raise ValueError(f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}") + + return input_ids, segment_ids, input_mask + if len(graph_inputs) != 3: raise ValueError("Expect the graph to have 3 inputs. Got {}".format(len(graph_inputs))) diff --git a/onnxruntime/python/tools/bert/compare_bert_results.py b/onnxruntime/python/tools/bert/compare_bert_results.py index 248a25d246..13b603a6a2 100644 --- a/onnxruntime/python/tools/bert/compare_bert_results.py +++ b/onnxruntime/python/tools/bert/compare_bert_results.py @@ -78,9 +78,11 @@ def compare(baseline_results, treatment_results, verbose, rtol=1e-3, atol=1e-4): def run_test(baseline_model, optimized_model, output_dir, batch_size, sequence_length, use_gpu, test_cases, seed, - use_openmp, verbose, rtol, atol): + use_openmp, verbose, rtol, atol, input_ids_name, segment_ids_name, input_mask_name): + # Try deduce input names from optimized model. - input_ids, segment_ids, input_mask = get_bert_inputs(optimized_model) + input_ids, segment_ids, input_mask = get_bert_inputs(optimized_model, input_ids_name, segment_ids_name, + input_mask_name) # Use random mask length for accuracy test. It might introduce slight inflation in latency reported in this script. all_inputs = generate_test_data(batch_size, @@ -161,6 +163,10 @@ def parse_arguments(): parser.add_argument('--verbose', required=False, action='store_true', help="print verbose information") parser.set_defaults(verbose=False) + parser.add_argument('--input_ids', required=False, type=str, default=None, help="input name for input ids") + parser.add_argument('--segment_ids', required=False, type=str, default=None, help="input name for segment ids") + parser.add_argument('--input_mask', required=False, type=str, default=None, help="input name for attention mask") + args = parser.parse_args() return args @@ -174,7 +180,8 @@ def main(): path.mkdir(parents=True, exist_ok=True) run_test(args.baseline_model, args.optimized_model, args.output_dir, args.batch_size, args.sequence_length, - args.use_gpu, args.samples, args.seed, args.openmp, args.verbose, args.rtol, args.atol) + args.use_gpu, args.samples, args.seed, args.openmp, args.verbose, args.rtol, args.atol, args.input_ids, + args.segment_ids, args.input_mask) if __name__ == "__main__": diff --git a/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb b/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb index aaff0445f9..d7f1bd6fdf 100644 --- a/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb +++ b/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb @@ -484,25 +484,25 @@ ] } ], - "source": [ - "import os\n", - "import wget\n", - "\n", - "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", - "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n", - "\n", - "script_dir = './bert_scripts'\n", - "if not os.path.exists(script_dir):\n", - " os.makedirs(script_dir)\n", - "\n", - "for filename in script_files:\n", - " target_file = os.path.join(script_dir, filename)\n", - " if enable_overwrite and os.path.exists(target_file):\n", - " os.remove(target_file)\n", - " if not os.path.exists(target_file):\n", - " wget.download(url_prfix + filename, target_file)\n", - " print(\"Downloaded\", filename)" - ] + "source": [ + "import os\n", + "import wget\n", + "\n", + "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", + "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n", + "\n", + "script_dir = './bert_scripts'\n", + "if not os.path.exists(script_dir):\n", + " os.makedirs(script_dir)\n", + "\n", + "for filename in script_files:\n", + " target_file = os.path.join(script_dir, filename)\n", + " if enable_overwrite and os.path.exists(target_file):\n", + " os.remove(target_file)\n", + " if not os.path.exists(target_file):\n", + " wget.download(url_prfix + filename, target_file)\n", + " print(\"Downloaded\", filename)" + ] }, { "cell_type": "markdown", diff --git a/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_GPU.ipynb b/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_GPU.ipynb index d9ac643786..a0a0e32e1f 100644 --- a/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_GPU.ipynb +++ b/onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_GPU.ipynb @@ -21,7 +21,7 @@ "source": [ "In this tutorial, you'll be introduced to how to load a Bert model from PyTorch, convert it to ONNX, and inference it for high performance using ONNX Runtime and NVIDIA GPU. In the following sections, we are going to use the Bert model trained with Stanford Question Answering Dataset (SQuAD) dataset as an example. Bert SQuAD model is used in question answering scenarios, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable.\n", "\n", - "This notebook is for CPU inference. For GPU inference, please look at another notebook [Inference PyTorch Bert Model with ONNX Runtime on CPU](PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb)." + "This notebook is for GPU inference. For CPU inference, please look at another notebook [Inference PyTorch Bert Model with ONNX Runtime on CPU](PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb)." ] }, { @@ -603,25 +603,25 @@ ] } ], - "source": [ - "import os\n", - "import wget\n", - "\n", - "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", - "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n", - "\n", - "script_dir = './bert_scripts'\n", - "if not os.path.exists(script_dir):\n", - " os.makedirs(script_dir)\n", - "\n", - "for filename in script_files:\n", - " target_file = os.path.join(script_dir, filename)\n", - " if enable_overwrite and os.path.exists(target_file):\n", - " os.remove(target_file)\n", - " if not os.path.exists(target_file):\n", - " wget.download(url_prfix + filename, target_file)\n", - " print(\"Downloaded\", filename)" - ] + "source": [ + "import os\n", + "import wget\n", + "\n", + "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", + "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n", + "\n", + "script_dir = './bert_scripts'\n", + "if not os.path.exists(script_dir):\n", + " os.makedirs(script_dir)\n", + "\n", + "for filename in script_files:\n", + " target_file = os.path.join(script_dir, filename)\n", + " if enable_overwrite and os.path.exists(target_file):\n", + " os.remove(target_file)\n", + " if not os.path.exists(target_file):\n", + " wget.download(url_prfix + filename, target_file)\n", + " print(\"Downloaded\", filename)" + ] }, { "cell_type": "markdown", diff --git a/onnxruntime/python/tools/bert/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb b/onnxruntime/python/tools/bert/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb index d1b0773cf7..3d084324f1 100644 --- a/onnxruntime/python/tools/bert/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb +++ b/onnxruntime/python/tools/bert/notebooks/Tensorflow_Keras_Bert-Squad_OnnxRuntime_CPU.ipynb @@ -100,30 +100,30 @@ ] } ], - "source": [ - "import os\n", - "import wget\n", - "\n", - "cache_dir = \"./squad\"\n", - "output_dir = \"./output\"\n", - "script_dir = './bert_scripts'\n", - "\n", - "for directory in [cache_dir, output_dir, script_dir]:\n", - " if not os.path.exists(directory):\n", - " os.makedirs(directory)\n", - "\n", - "# Download scripts for BERT optimization.\n", - "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", - "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py']\n", - "\n", - "for filename in script_files:\n", - " target_file = os.path.join(script_dir, filename)\n", - " if enable_overwrite and os.path.exists(target_file):\n", - " os.remove(target_file)\n", - " if not os.path.exists(target_file):\n", - " wget.download(url_prfix + filename, target_file)\n", - " print(\"Downloaded\", filename)" - ] + "source": [ + "import os\n", + "import wget\n", + "\n", + "cache_dir = \"./squad\"\n", + "output_dir = \"./output\"\n", + "script_dir = './bert_scripts'\n", + "\n", + "for directory in [cache_dir, output_dir, script_dir]:\n", + " if not os.path.exists(directory):\n", + " os.makedirs(directory)\n", + "\n", + "# Download scripts for BERT optimization.\n", + "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n", + "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py']\n", + "\n", + "for filename in script_files:\n", + " target_file = os.path.join(script_dir, filename)\n", + " if enable_overwrite and os.path.exists(target_file):\n", + " os.remove(target_file)\n", + " if not os.path.exists(target_file):\n", + " wget.download(url_prfix + filename, target_file)\n", + " print(\"Downloaded\", filename)" + ] }, { "cell_type": "markdown", diff --git a/onnxruntime/python/tools/bert/test_bert_optimization.py b/onnxruntime/python/tools/bert/test_bert_optimization.py index 6fa0422c32..528acf85e5 100644 --- a/onnxruntime/python/tools/bert/test_bert_optimization.py +++ b/onnxruntime/python/tools/bert/test_bert_optimization.py @@ -29,7 +29,9 @@ BERT_TEST_MODELS = { "bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx', "bert_keras_squad": - 'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx' + 'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx', + "gpt2": + 'test_data\\gpt2_pytorch1.4_opset11_no_past\\GPT2Model.onnx' } @@ -202,8 +204,8 @@ class TestBertOptimization(unittest.TestCase): 'Attention': 12, 'LayerNormalization': 0, 'SkipLayerNormalization': 24, - 'BiasGelu': 0, - 'Gelu': 12, + 'BiasGelu': 12, + 'Gelu': 0, 'FastGelu': 0 } self.verify_node_count(bert_model, expected_node_count) @@ -222,6 +224,28 @@ class TestBertOptimization(unittest.TestCase): self.assertTrue(bert_model.is_fully_optimized()) + def test_gpt2(self): + input = BERT_TEST_MODELS['gpt2'] + bert_model = optimize_model(input, + 'gpt2', + gpu_only=False, + num_heads=2, + hidden_size=4, + sequence_length=2, + input_int32=False, + float16=False) + + expected_node_count = { + 'EmbedLayerNormalization': 0, + 'Attention': 12, + 'Gelu': 0, + 'FastGelu': 12, + 'BiasGelu': 0, + 'LayerNormalization': 25, + 'SkipLayerNormalization': 0 + } + self.verify_node_count(bert_model, expected_node_count) + if __name__ == '__main__': unittest.main() diff --git a/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/GPT2Model.onnx b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/GPT2Model.onnx new file mode 100644 index 0000000000..c1b13ef70b Binary files /dev/null and b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/GPT2Model.onnx differ diff --git a/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/generate_tiny_gpt2_model.py b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/generate_tiny_gpt2_model.py new file mode 100644 index 0000000000..9a33a872e5 --- /dev/null +++ b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/generate_tiny_gpt2_model.py @@ -0,0 +1,394 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- +# This tool generates a tiny GPT2 model for testing fusion script. +# You can use benchmark_gpt2.py to get a gpt2 ONNX model as input of this tool. + +import onnx +import onnx.utils +import sys +import argparse +import numpy as np +from onnx import ModelProto, TensorProto, numpy_helper +from OnnxModel import OnnxModel +import os +import onnxruntime +import random +from pathlib import Path +import timeit + +DICT_SIZE = 20 +SEQ_LEN = 2 +""" This class creates a tiny bert model for test purpose. """ + +# parameters of input base model. +old_parameters = { + "seq_len": 5, + "hidden_size": 768, + "num_heads": 12, + "size_per_head": 64, + "word_dict_size": [50257], # list of supported dictionary size. + "max_word_position": 1024 +} + +# parameters of output tiny model. +new_parameters = { + "seq_len": SEQ_LEN, + "hidden_size": 4, + "num_heads": 2, + "size_per_head": 2, + "word_dict_size": DICT_SIZE, + "max_word_position": 8 +} + + +class TinyBertOnnxModel(OnnxModel): + + def __init__(self, model): + super(TinyBertOnnxModel, self).__init__(model) + self.resize_model() + + def resize_weight(self, initializer_name, target_shape): + weight = self.get_initializer(initializer_name) + w = numpy_helper.to_array(weight) + + target_w = w + if len(target_shape) == 1: + target_w = w[:target_shape[0]] + elif len(target_shape) == 2: + target_w = w[:target_shape[0], :target_shape[1]] + elif len(target_shape) == 3: + target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]] + elif len(target_shape) == 4: + target_w = w[:target_shape[0], :target_shape[1], :target_shape[2], :target_shape[3]] + else: + print("at most 3 dimensions") + + tensor = onnx.helper.make_tensor(name=initializer_name + '_resize', + data_type=TensorProto.FLOAT, + dims=target_shape, + vals=target_w.flatten().tolist()) + + return tensor + + def resize_model(self): + graph = self.model.graph + initializers = graph.initializer + + for input in graph.input: + if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]): + print("input", input.name, input.type.tensor_type.shape) + input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"] + print("=>", input.type.tensor_type.shape) + + reshapes = {} + for initializer in initializers: + tensor = numpy_helper.to_array(initializer) + if initializer.data_type == TensorProto.FLOAT: + dtype = np.float32 + elif initializer.data_type == TensorProto.INT32: + dtype = np.int32 + elif initializer.data_type == TensorProto.INT64: + dtype = np.int64 + else: + print("data type not supported by this tool:", dtype) + + if len(tensor.shape) == 1 and tensor.shape[0] == 1: + if tensor == old_parameters["num_heads"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["seq_len"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name)) + elif tensor == old_parameters["size_per_head"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif tensor == 4 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif tensor == 3 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 3 * old_parameters["hidden_size"], "=>[", 3 * new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([3 * new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif len(tensor.shape) == 0: + if tensor == old_parameters["num_heads"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["num_heads"], "=>", new_parameters["num_heads"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name)) + elif tensor == old_parameters["seq_len"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["seq_len"], "=>", new_parameters["seq_len"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name)) + elif tensor == old_parameters["size_per_head"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["size_per_head"], "=>", new_parameters["size_per_head"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["hidden_size"], "=>", new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 4 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 3 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 3 * old_parameters["hidden_size"], "=>", 3 * new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(3 * new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]): + print("initializer type={}".format(initializer.data_type), initializer.name, + 1.0 / np.sqrt(old_parameters["size_per_head"]), "=>", + 1.0 / np.sqrt(new_parameters["size_per_head"])) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype), + initializer.name)) + elif tensor == np.sqrt(old_parameters["size_per_head"]): + print("initializer type={}".format(initializer.data_type), initializer.name, + np.sqrt(old_parameters["size_per_head"]), "=>", np.sqrt(new_parameters["size_per_head"])) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(np.sqrt(new_parameters["size_per_head"]), dtype=dtype), + initializer.name)) + + new_shape = [] + shape_changed = False + for dim in tensor.shape: + if (dim == old_parameters["hidden_size"]): + new_shape.append(new_parameters["hidden_size"]) + shape_changed = True + elif (dim == 4 * old_parameters["hidden_size"]): + new_shape.append(4 * new_parameters["hidden_size"]) + shape_changed = True + elif (dim == 3 * old_parameters["hidden_size"]): + new_shape.append(3 * new_parameters["hidden_size"]) + shape_changed = True + elif (dim in old_parameters["word_dict_size"]): + new_shape.append(new_parameters["word_dict_size"]) + shape_changed = True + elif (dim == old_parameters["max_word_position"]): + new_shape.append(new_parameters["max_word_position"]) + shape_changed = True + else: + new_shape.append(dim) + if shape_changed: + reshapes[initializer.name] = new_shape + print("initializer", initializer.name, tensor.shape, "=>", new_shape) + + for initializer_name in reshapes: + self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize') + tensor = self.resize_weight(initializer_name, reshapes[initializer_name]) + self.model.graph.initializer.extend([tensor]) + + # Add node name, replace split node attribute. + nodes_to_add = [] + nodes_to_remove = [] + for i, node in enumerate(graph.node): + if node.op_type == "Split": + nodes_to_add.append( + onnx.helper.make_node('Split', + node.input, + node.output, + name="Split_{}".format(i), + axis=2, + split=[ + new_parameters["hidden_size"], new_parameters["hidden_size"], + new_parameters["hidden_size"] + ])) + nodes_to_remove.append(node) + print("update split", + [new_parameters["hidden_size"], new_parameters["hidden_size"], new_parameters["hidden_size"]]) + if node.op_type == "Constant": + for att in node.attribute: + if att.name == 'value': + if numpy_helper.to_array(att.t) == old_parameters["num_heads"]: + nodes_to_add.append( + onnx.helper.make_node('Constant', + inputs=node.input, + outputs=node.output, + value=onnx.helper.make_tensor(name=att.t.name, + data_type=TensorProto.INT64, + dims=[], + vals=[new_parameters["num_heads"] + ]))) + print("constant", att.t.name, old_parameters["num_heads"], "=>", + new_parameters["num_heads"]) + if numpy_helper.to_array(att.t) == np.sqrt(old_parameters["size_per_head"]): + nodes_to_add.append( + onnx.helper.make_node('Constant', + inputs=node.input, + outputs=node.output, + value=onnx.helper.make_tensor( + name=att.t.name, + data_type=TensorProto.FLOAT, + dims=[], + vals=[np.sqrt(new_parameters["size_per_head"])]))) + print("constant", att.t.name, np.sqrt(old_parameters["size_per_head"]), "=>", + np.sqrt(new_parameters["size_per_head"])) + else: + node.name = node.op_type + "_" + str(i) + for node in nodes_to_remove: + graph.node.remove(node) + graph.node.extend(nodes_to_add) + + def remove_past_outputs(self): + keep_output_names = [self.model.graph.output[0].name] # remove past state outputs which is not needed. + print(f"Prune graph to keep the first output and drop past state outputs:{keep_output_names}") + self.prune_graph(keep_output_names) + + +def generate_test_data(onnx_file, + output_path, + batch_size, + sequence_length, + use_cpu=True, + input_tensor_only=False, + dictionary_size=DICT_SIZE, + test_cases=1, + output_optimized_model=False): + + input_data_type = np.int64 + for test_case in range(test_cases): + input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=input_data_type) + tensor_1 = numpy_helper.from_array(input_1, 'input_ids') + + path = os.path.join(output_path, 'test_data_set_' + str(test_case)) + try: + os.mkdir(path) + except OSError: + print("Creation of the directory %s failed" % path) + else: + print("Successfully created the directory %s " % path) + + if input_tensor_only: + return + + sess_options = onnxruntime.SessionOptions() + sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider']) + + input1_name = sess.get_inputs()[0].name + output_names = [output.name for output in sess.get_outputs()] + inputs = {input1_name: input_1} + result = sess.run(output_names, inputs) + + with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f: + f.write(tensor_1.SerializeToString()) + + for i, output_name in enumerate(output_names): + if i == 0: + tensor_result = numpy_helper.from_array( + np.asarray(result[i]).reshape((batch_size, sequence_length, new_parameters["hidden_size"])), + output_names[i]) + with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f: + f.write(tensor_result.SerializeToString()) + else: + tensor_result = numpy_helper.from_array( + np.asarray(result[i]).reshape( + (2, batch_size, new_parameters["num_heads"], sequence_length, new_parameters["size_per_head"])), + output_names[i]) + with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f: + f.write(tensor_result.SerializeToString()) + + start_time = timeit.default_timer() + + sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED + + if output_optimized_model: + path_prefix = onnx_file[:-5] #remove .onnx suffix + if use_cpu: + sess_options.optimized_model_filepath = path_prefix + "_optimized_cpu.onnx" + else: + sess_options.optimized_model_filepath = path_prefix + "_optimized_gpu.onnx" + + session = onnxruntime.InferenceSession(onnx_file, sess_options) + if use_cpu: + session.set_providers(['CPUExecutionProvider']) # use cpu + else: + if 'CUDAExecutionProvider' not in session.get_providers(): + print("Warning: GPU not found") + continue + outputs = session.run(None, inputs) + evalTime = timeit.default_timer() - start_time + if not np.allclose(outputs[0], result[0], rtol=1e-04, atol=1e-05): + print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format( + use_cpu, result[0].tolist(), outputs[0].tolist())) + print("** Evaluation done in total {} secs".format(evalTime)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--input', required=True, type=str) + parser.add_argument('--output', required=True, type=str) + parser.add_argument('--float16', required=False, action='store_true') + parser.set_defaults(float16=False) + parser.add_argument('--no_past_outputs', required=False, action='store_true') + parser.set_defaults(no_past_outputs=False) + parser.add_argument('--output_optimized_model', required=False, action='store_true') + parser.set_defaults(output_optimized_model=False) + args = parser.parse_args() + + model = ModelProto() + with open(args.input, "rb") as f: + model.ParseFromString(f.read()) + + bert_model = TinyBertOnnxModel(model) + + if args.float16: + bert_model.convert_model_float32_to_float16() + + if args.no_past_outputs: + bert_model.remove_past_outputs() + + bert_model.update_graph() + bert_model.remove_unused_constant() + + print("opset verion", bert_model.model.opset_import[0].version) + + with open(args.output, "wb") as out: + out.write(bert_model.model.SerializeToString()) + + p = Path(args.output) + data_path = p.parent + + batch_size = 1 + sequence_length = SEQ_LEN + + generate_test_data(args.output, + data_path, + batch_size, + sequence_length, + use_cpu=not args.float16, + output_optimized_model=args.output_optimized_model) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/input_0.pb b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/input_0.pb new file mode 100644 index 0000000000..42d734d9ae Binary files /dev/null and b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/input_0.pb differ diff --git a/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/output_0.pb b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/output_0.pb new file mode 100644 index 0000000000..0147151612 --- /dev/null +++ b/onnxruntime/python/tools/bert/test_data/gpt2_pytorch1.4_opset11_no_past/test_data_set_0/output_0.pb @@ -0,0 +1 @@ +B hidden_statesJ ÙaÊ>ÛÑ>&ÏIÀ¢?½;g>v,²>©3CÀYŽ²? \ No newline at end of file