diff --git a/onnxruntime/python/tools/bert/BertOnnxModel.py b/onnxruntime/python/tools/bert/BertOnnxModel.py index f30aeb56c8..cb4b0b60b6 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModel.py +++ b/onnxruntime/python/tools/bert/BertOnnxModel.py @@ -15,12 +15,12 @@ from OnnxModel import OnnxModel logger = logging.getLogger(__name__) class BertOnnxModel(OnnxModel): - def __init__(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only, verbose): + def __init__(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only): assert num_heads > 0 assert hidden_size % num_heads == 0 assert sequence_length > 0 - super(BertOnnxModel, self).__init__(model, verbose) + super(BertOnnxModel, self).__init__(model) self.num_heads = num_heads self.sequence_length = sequence_length self.hidden_size = hidden_size @@ -36,13 +36,6 @@ class BertOnnxModel(OnnxModel): self.bert_inputs = [] - # constant node names - self.normalize_name = "SkipLayerNormalization" - self.attention_name = 'Attention' - - def get_normalize_nodes(self): - return self.get_nodes_by_op_type(self.normalize_name) - def normalize_children_types(self): return ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization'] @@ -100,7 +93,7 @@ class BertOnnxModel(OnnxModel): q_bias = self.get_initializer(q_add.input[1]) k_bias = self.get_initializer(k_add.input[1]) v_bias = self.get_initializer(v_add.input[1]) - + qw = numpy_helper.to_array(q_weight) assert qw.shape == (self.hidden_size, self.hidden_size) @@ -123,7 +116,7 @@ class BertOnnxModel(OnnxModel): qkv_bias = np.stack((qb, kb, vb), axis=-2) - attention_node_name = self.create_node_name(self.attention_name) + attention_node_name = self.create_node_name('Attention') weight = onnx.helper.make_tensor(name=attention_node_name + '_qkv_weight', data_type=TensorProto.FLOAT, @@ -143,8 +136,10 @@ class BertOnnxModel(OnnxModel): 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(self.attention_name, - inputs=[input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias', mask_index], outputs=[output], + attention_node = onnx.helper.make_node( + 'Attention', + inputs=[input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias', mask_index], + outputs=[output], name=attention_node_name) attention_node.domain = "com.microsoft" attention_node.attribute.extend([onnx.helper.make_attribute("num_heads", self.num_heads)]) @@ -152,13 +147,17 @@ class BertOnnxModel(OnnxModel): self.add_node(attention_node) def fuse_attention(self): + """ + Fuse Attention subgraph into one Attention node. + """ 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_normalize_nodes(): + skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + 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 @@ -180,6 +179,7 @@ class BertOnnxModel(OnnxModel): v_nodes = self.match_parent_path(matmul_qkv, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0]) if v_nodes is None: + logger.debug("fuse_attention: failed to match v path") continue (transpose_v, reshape_v, add_v, matmul_v) = v_nodes @@ -187,11 +187,13 @@ class BertOnnxModel(OnnxModel): if qk_nodes is None: qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, 0, 0]) if qk_nodes is None: - continue + logger.debug("fuse_attention: failed to match qk path") + continue (softmax_qk, add_qk, div_qk, matmul_qk) = qk_nodes q_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, 0]) if q_nodes is None: + logger.debug("fuse_attention: failed to match q path") continue (transpose_q, reshape_q, add_q, matmul_q) = q_nodes @@ -199,6 +201,7 @@ class BertOnnxModel(OnnxModel): if k_nodes is None: k_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0, 0]) if k_nodes is None: + logger.debug("fuse_attention: failed to match k path") continue (transpose_k, transpose_k_2, reshape_k, add_k, matmul_k) = k_nodes else: @@ -206,6 +209,7 @@ class BertOnnxModel(OnnxModel): mask_nodes = self.match_parent_path(add_qk, ['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]) if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") continue (mul_mask, sub_mask, cast_mask, unsqueeze_mask, unsqueeze_mask_0) = mask_nodes @@ -451,13 +455,20 @@ class BertOnnxModel(OnnxModel): nodes_to_remove = [] nodes_to_add = [] - for node in self.get_normalize_nodes(): + skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + for node in skip_layer_norm_nodes: if len(node.input) != 4: continue - nodes = self.match_parent_path(node, ['Add', 'MatMul'], [0, None]) + return_indice = [] + nodes = self.match_parent_path(node, ['Add', 'MatMul'], [None, None], None, return_indice) if nodes is None: continue + assert len(return_indice) == 2 + add_input_index = return_indice[0] + if add_input_index >= 2: + continue + (add, matmul) = nodes # bias should be one dimension @@ -470,19 +481,22 @@ class BertOnnxModel(OnnxModel): bias_weight = numpy_helper.to_array(initializer) break if bias_weight is None: + logger.debug(f"Bias weight not found") continue if len(bias_weight.shape) != 1: + logger.debug(f"Bias weight is not 1D") continue subgraph_nodes = [node, add] if not self.is_safe_to_fuse_nodes(subgraph_nodes, [node.output[0]], input_name_to_nodes, output_name_to_node): + logger.debug(f"Skip fusing SkipLayerNormalization with Bias since it is not safe") continue nodes_to_remove.extend(subgraph_nodes) - new_node = onnx.helper.make_node(self.normalize_name, - inputs=[matmul.output[0], node.input[1], node.input[2], node.input[3], add.input[bias_index]], + new_node = onnx.helper.make_node("SkipLayerNormalization", + inputs=[node.input[1 - add_input_index], matmul.output[0], node.input[2], node.input[3], add.input[bias_index]], outputs=node.output, - name=self.create_node_name(self.normalize_name, self.normalize_name + "_AddBias_")) + name=self.create_node_name("SkipLayerNormalization", "SkipLayerNorm_AddBias_")) new_node.domain = "com.microsoft" nodes_to_add.append(new_node) @@ -640,7 +654,9 @@ class BertOnnxModel(OnnxModel): # Find the first normalize node could be embedding layer. normalize_node = None - for node in self.get_normalize_nodes(): + + skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + for node in skip_layer_norm_nodes: if self.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is not None: if self.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) is not None: normalize_node = node @@ -787,15 +803,16 @@ class BertOnnxModel(OnnxModel): self.model = onnx.helper.make_model(graph_def, producer_name='bert model optimizer') if isinstance(batch_size, str): - self.update_dynamic_batch_io(batch_size) + self.use_dynamic_axes(batch_size, None) # restore opset version self.model.opset_import[0].version = original_opset_version - # Update input and output using dynamic batch - def update_dynamic_batch_io(self, dynamic_batch_dim='batch'): - + def use_dynamic_axes(self, dynamic_batch_dim='batch_size', dynamic_seq_len='max_seq_len'): + """ + Update input and output shape to use dynamic axes. + """ bert_inputs = self.get_bert_inputs() dynamic_batch_inputs = {} for input in self.model.graph.input: @@ -803,34 +820,36 @@ class BertOnnxModel(OnnxModel): if bert_input == input.name: dim_proto = input.type.tensor_type.shape.dim[0] dim_proto.dim_param = dynamic_batch_dim + if dynamic_seq_len is not None: + dim_proto = input.type.tensor_type.shape.dim[1] + dim_proto.dim_param = dynamic_seq_len for output in self.model.graph.output: dim_proto = output.type.tensor_type.shape.dim[0] dim_proto.dim_param = dynamic_batch_dim - """ - Layer Normalization will fuse Add + LayerNormalization into one node: - +----------------------+ - | | - | v - Add --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add - (axis=2 or -1) | (Y=2) (axis=2 or -1) (E-6 or E-12 or 0) ^ - | | - +-----------------------------------------------+ - - It also handles cases of duplicated sub nodes exported from older version of PyTorch: - +----------------------+ - | v - | +-------> Sub-----------------------------------------------+ - | | | - | | v - Add --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add - | ^ - | | - +----------------------+ - - """ def fuse_layer_norm(self): + """ + Fuse Layer Normalization subgraph into one node LayerNormalization: + +----------------------+ + | | + | v + [Root] --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add + (axis=2 or -1) | (Y=2) (axis=2 or -1) (E-6 or E-12 or 0) ^ + | | + +-----------------------------------------------+ + + It also handles cases of duplicated sub nodes exported from older version of PyTorch: + +----------------------+ + | v + | +-------> Sub-----------------------------------------------+ + | | | + | | v + [Root] --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add + | ^ + | | + +----------------------+ + """ input_name_to_nodes = self.input_name_to_nodes() output_name_to_node = self.output_name_to_node() @@ -896,32 +915,51 @@ class BertOnnxModel(OnnxModel): 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 parent.op_type == 'Add' and self.is_safe_to_fuse_nodes([parent] + subgraph_nodes, last_add_node.output, input_name_to_nodes, output_name_to_node): - nodes_to_remove.append(parent) - normalize_node = onnx.helper.make_node(self.normalize_name, - inputs=[parent.input[0], parent.input[1], weight_input, bias_input], - outputs=[last_add_node.output[0]], - name=self.create_node_name(self.normalize_name, name_prefix="SkipLayerNorm")) + normalize_node = onnx.helper.make_node('LayerNormalization', + inputs=[node.input[0], weight_input, bias_input], + outputs=[last_add_node.output[0]]) + normalize_node.attribute.extend([onnx.helper.make_attribute("epsilon", add_weight)]) + layernorm_nodes.extend([normalize_node]) + + self.remove_nodes(nodes_to_remove) + self.add_nodes(layernorm_nodes) + logger.info(f"Fused LayerNormalization count: {len(layernorm_nodes)}") + + def fuse_skip_layer_norm(self): + """ + Fuse Add + LayerNormalization into one node: SkipLayerNormalization + """ + input_name_to_nodes = self.input_name_to_nodes() + output_name_to_node = self.output_name_to_node() + + nodes_to_remove = [] + skip_layernorm_nodes = [] + for node in self.nodes(): + if node.op_type == 'LayerNormalization': + add = self.get_parent(node, 0, output_name_to_node) + if add is None: + continue + + if add.op_type == 'Add' and self.is_safe_to_fuse_nodes([add, node], node.output, input_name_to_nodes, output_name_to_node): + nodes_to_remove.extend([add, node]) + normalize_node = onnx.helper.make_node("SkipLayerNormalization", + inputs=[add.input[0], add.input[1], node.input[1], node.input[2]], + outputs=[node.output[0]], + name=self.create_node_name("SkipLayerNormalization", name_prefix="SkipLayerNorm")) normalize_node.domain = "com.microsoft" skip_layernorm_nodes.extend([normalize_node]) - else: - normalize_node = onnx.helper.make_node('LayerNormalization', - inputs=[node.input[0], weight_input, bias_input], - outputs=[last_add_node.output[0]]) - normalize_node.attribute.extend([onnx.helper.make_attribute("epsilon", add_weight)]) - layernorm_nodes.extend([normalize_node]) self.remove_nodes(nodes_to_remove) self.add_nodes(skip_layernorm_nodes) - self.add_nodes(layernorm_nodes) logger.info(f"Fused SkipLayerNormalization count: {len(skip_layernorm_nodes)}") - logger.info(f"Fused LayerNormalization count: {len(layernorm_nodes)}") def preprocess(self): return + def postprocess(self): + return + def optimize(self): - self. preprocess() self.fuse_layer_norm() # FastGelu uses approximation for Gelu. It is faster. @@ -929,22 +967,28 @@ class BertOnnxModel(OnnxModel): gelu_op_name = 'FastGelu' if use_approximation else 'Gelu' self.fuse_gelu(gelu_op_name) + self.preprocess() + self.fuse_reshape() + self.fuse_skip_layer_norm() 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() + self.postprocess() + if self.float16: self.convert_model_float32_to_float16() self.remove_unused_constant() # Use symbolic batch dimension in input and output. - self.update_dynamic_batch_io() + self.use_dynamic_axes() logger.info(f"opset verion: {self.model.opset_import[0].version}") diff --git a/onnxruntime/python/tools/bert/BertOnnxModelKeras.py b/onnxruntime/python/tools/bert/BertOnnxModelKeras.py new file mode 100644 index 0000000000..af76d8a021 --- /dev/null +++ b/onnxruntime/python/tools/bert/BertOnnxModelKeras.py @@ -0,0 +1,265 @@ +#------------------------------------------------------------------------- +# 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 BertOnnxModelTF import BertOnnxModelTF + +logger = logging.getLogger(__name__) + +class BertOnnxModelKeras(BertOnnxModelTF): + def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only): + super().__init__(model, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only) + + def fuse_attention(self): + input_name_to_nodes = self.input_name_to_nodes() + output_name_to_node = self.output_name_to_node() + + nodes_to_remove = [] + attention_count = 0 + + skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + for normalize_node in skip_layer_norm_nodes: + # SkipLayerNormalization has two inputs, and one of them is the root input for attention. + parent = self.get_parent(normalize_node, 0) + if parent is None or parent.op_type not in ["SkipLayerNormalization", "EmbedLayerNormalization"]: + if parent.op_type == 'Add': + parent = self.get_parent(normalize_node, 1) + if parent is None or parent.op_type not in ["SkipLayerNormalization", "EmbedLayerNormalization"]: + logger.debug("First input for skiplayernorm: {}".format(parent.op_type if parent is not None else None)) + continue + else: + logger.debug("First input for skiplayernorm: {}".format(parent.op_type if parent is not None else None)) + continue + else: + # TODO: shall we add back the checking of children op types. + pass + + qkv_nodes = self.match_parent_path(normalize_node, + ['Add', 'Reshape', 'MatMul', 'Reshape', 'Transpose', 'MatMul'], + [ None, 0, 0, 0, 0, 0]) + if qkv_nodes is None: + logger.debug("Failed to match qkv nodes") + continue + (add, extra_reshape_0, matmul, reshape_qkv, transpose_qkv, matmul_qkv) = qkv_nodes + logger.debug("Matched qkv nodes") + + v_nodes = self.match_parent_path(matmul_qkv, + ['Transpose', 'Reshape', 'Add', 'Reshape', 'MatMul'], + [ 1, 0, 0, 0, 0]) + if v_nodes is None: + logger.debug("Failed to match v path") + continue + (transpose_v, reshape_v, add_v, extra_reshape_1, matmul_v) = v_nodes + + qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'MatMul'], [0, 0, 0]) + if qk_nodes is not None: + (softmax_qk, sub_qk, matmul_qk) = qk_nodes + q_nodes = self.match_parent_path(matmul_qk, + ['Mul', 'Transpose', 'Reshape', 'Add', 'Reshape', 'MatMul'], + [ 0, None, 0, 0, 0, 0]) + if q_nodes is not None: + (mul_q, transpose_q, reshape_q, add_q, extra_reshape_2, matmul_q) = q_nodes + + else: + qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, 0, None]) + if qk_nodes is None: + logger.debug("Failed to match qk path") + continue + (softmax_qk, add_qk, mul_qk, matmul_qk) = qk_nodes + + q_nodes = self.match_parent_path(matmul_qk, + ['Transpose', 'Reshape', 'Add', 'Reshape', 'MatMul'], + [ 0, 0, 0, 0, 0]) + if q_nodes is not None: + (transpose_q, reshape_q, add_q, extra_reshape_2, matmul_q) = q_nodes + + if q_nodes is None: + logger.debug("Failed to match q path") + continue + + k_nodes = self.match_parent_path(matmul_qk, + ['Transpose', 'Reshape', 'Add', 'Reshape', 'MatMul'], + [ 1, 0, 0, 0, 0]) + if k_nodes is None: + logger.debug("Failed to match k path") + continue + (transpose_k, reshape_k, add_k, extra_reshape_3, matmul_k) = k_nodes + + + mask_nodes = self.match_parent_path(qk_nodes[1], + ['Mul', 'Sub', 'Reshape', 'Cast'], + [ 1, None, 1, 0]) + if mask_nodes is None: + mask_nodes = self.match_parent_path(qk_nodes[1], + ['Mul', 'Sub', 'Cast', 'Slice', 'Unsqueeze'], + [ 1, 1, 1, 0, 0]) + if mask_nodes is None: + logger.debug("Failed to match mask path") + continue + (mul_mask, sub_mask, cast_mask, slice_mask, unsqueeze_mask) = mask_nodes + else: + (mul_mask, sub_mask, reshape_mask, cast_mask) = mask_nodes + + if not self.has_constant_input(sub_mask, 1): + logger.debug("Sub node expected to have an input with constant value 1.0.") + continue + + root_input = matmul_v.input[0] + root_node = output_name_to_node[root_input] + is_same_root = (root_node == parent or (root_node.op_type == 'Reshape' and root_node.input[0] == parent.output[0])) + if matmul_q.input[0] == root_input and matmul_v.input[0] == root_input and is_same_root: + mask_index = self.process_mask(mask_nodes[-1].input[0]) + logger.debug("Create an Attention node.") + self.create_attention_node(mask_index, matmul_q, matmul_k, matmul_v, add_q, add_k, add_v, parent.output[0], reshape_qkv.output[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) + if root_node.op_type == 'Reshape': + nodes_to_remove.append(root_node) + attention_count += 1 + + nodes_to_remove.append(extra_reshape_0) + self.replace_node_input(add, extra_reshape_0.output[0], matmul.output[0]) + else: + logger.debug("Root node not matched.") + continue + self.remove_nodes(nodes_to_remove) + self.update_graph() + logger.info(f"Fused Attention count:{attention_count}") + + def fuse_embedding(self, node, output_name_to_node): + assert node.op_type == 'LayerNormalization' + pos_embed_path2 = self.match_parent_path( + node, + ['Add', 'Add', 'Gather'], + [ 0, 0, 0], + output_name_to_node) + if pos_embed_path2 is None: + logger.debug("failed to match pos_embed_path") + return False + + skip_node, add_node, gather_node = pos_embed_path2 + pos_initializer = self.get_initializer(add_node.input[1]) + if pos_initializer is None: + return False + + temp = numpy_helper.to_array(pos_initializer) + if len(temp.shape) == 3 and temp.shape[0] == 1: + tensor = numpy_helper.from_array(temp.reshape((temp.shape[1],temp.shape[2])), "position_embedding") + self.add_initializer(tensor) + logger.info("Found position embedding. name:{}, shape:{}".format(pos_initializer.name, temp.shape[1:])) + position_embedding = "position_embedding" + else: + logger.info("Failed to find position embedding. name:{}, shape:{}".format(pos_initializer.name, temp.shape)) + return False + + + word_initializer = self.get_initializer(gather_node.input[0]) + if word_initializer is None: + return False + + temp = numpy_helper.to_array(word_initializer) + if len(temp.shape) == 2: + logger.info("Found word embedding. name:{}, shape:{}".format(word_initializer.name, temp.shape)) + word_embedding = word_initializer.name + else: + logger.info("Failed to find word embedding. name:{}, shape:{}".format(word_initializer.name, temp.shape)) + return False + + gather = self.get_parent(skip_node, 1, output_name_to_node) + if gather is None or gather.op_type != "Gather": + return False + + segment_initializer = self.get_initializer(gather.input[0]) + if segment_initializer is None: + return False + + temp = numpy_helper.to_array(segment_initializer) + if len(temp.shape) == 2: + logger.info("Found segment embedding. name:{}, shape:{}".format(segment_initializer.name, temp.shape)) + segment_embedding = segment_initializer.name + else: + logger.info("Failed to find segment embedding. name:{}, shape:{}".format(segment_initializer.name, temp.shape)) + return False + + logger.info("Create Embedding node") + self.create_embedding_subgraph(node, word_embedding, segment_embedding, position_embedding) + return True + + def process_embedding(self): + """ + Automatically detect word, segment and position embeddings. + """ + logger.info("start processing embedding layer...") + output_name_to_node = self.output_name_to_node() + for node in self.nodes(): + if node.op_type == 'LayerNormalization': + if self.fuse_embedding(node, output_name_to_node): + return + + def remove_extra_reshape(self): + skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + reshape_removed = 0 + for skiplayernorm_node in skiplayernorm_nodes: + path = self.match_parent_path( + skiplayernorm_node, + ['Add', 'Reshape', 'MatMul', 'Reshape', 'Gelu', 'Add', 'Reshape', 'MatMul', 'SkipLayerNormalization'], + [ 0, 0, 0, 0, 0, 0, 0, 0, 0]) + if path is None: + continue + + add_1, reshape_1, matmul_1, reshape_2, gelu, add_2, reshape_3, matmul_2, skiplayernorm = path + add_2.input[0] = matmul_2.output[0] + self.remove_node(reshape_3) + matmul_1.input[0] = gelu.output[0] + self.remove_node(reshape_2) + add_1.input[0] = matmul_1.output[0] + self.remove_node(reshape_1) + reshape_removed += 3 + + logger.info(f"Remove {reshape_removed} Reshape nodes.") + + def remove_extra_reshape_2(self): + skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization") + reshape_removed = 0 + for skiplayernorm_node in skiplayernorm_nodes: + path = self.match_parent_path( + skiplayernorm_node, + ['Add', 'Reshape', 'MatMul', 'Reshape', 'Gelu', 'Add', 'Reshape', 'MatMul', 'Reshape', 'SkipLayerNormalization'], + [ None, 0, 0, 0, 0, 0, 0, 0, 0, 0]) + if path is None: + continue + + add_1, reshape_1, matmul_1, reshape_2, gelu, add_2, reshape_3, matmul_2, reshape_4, skiplayernorm = path + + matmul_2.input[0] = skiplayernorm.output[0] + self.remove_node(reshape_4) + + add_2.input[0] = matmul_2.output[0] + self.remove_node(reshape_3) + + matmul_1.input[0] = gelu.output[0] + self.remove_node(reshape_2) + + add_1.input[0] = matmul_1.output[0] + self.remove_node(reshape_1) + + reshape_removed += 4 + + logger.info(f"Remove {reshape_removed} Reshape nodes.") + + def postprocess(self): + self.remove_extra_reshape() + self.remove_extra_reshape_2() + self.prune_graph() diff --git a/onnxruntime/python/tools/bert/BertOnnxModelTF.py b/onnxruntime/python/tools/bert/BertOnnxModelTF.py index 57ad6520e7..c255ff6d52 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModelTF.py +++ b/onnxruntime/python/tools/bert/BertOnnxModelTF.py @@ -15,8 +15,8 @@ from BertOnnxModel import BertOnnxModel logger = logging.getLogger(__name__) class BertOnnxModelTF(BertOnnxModel): - def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only, verbose): - super().__init__(model, num_heads, hidden_size, sequence_length, verbose) + def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only): + super().__init__(model, num_heads, hidden_size, sequence_length) """ Fuse Gelu with Erf into one node: @@ -198,124 +198,6 @@ class BertOnnxModelTF(BertOnnxModel): self.remove_nodes(nodes_to_remove) self.add_nodes(nodes_to_add) - def __fuse_reshape_after_qkv(self, reshape_node, nodes_to_remove, nodes_to_add): - path0 = self.match_parent_path(reshape_node, ['Cast', 'Concat'], [1, 0]) - if path0 is None: - return - cast_node, concat_node = path0 - - if not len(concat_node.input) == 4: - return - - shape = [0,0] - - if self.get_initializer(concat_node.input[2]) and self.get_initializer(concat_node.input[3]): - concat_2 = self.get_initializer(concat_node.input[2]) - concat_3 = self.get_initializer(concat_node.input[3]) - shape.extend(numpy_helper.to_array(concat_2)) - shape.extend(numpy_helper.to_array(concat_3)) - else: - return - - shape_value = np.asarray(shape, dtype=np.int64) - - constant_shape_name = self.create_node_name('Constant', 'constant_shape') - new_node = onnx.helper.make_node('Constant', - inputs=[], - outputs=[constant_shape_name], - value=onnx.helper.make_tensor(name='const_tensor', - data_type=TensorProto.INT64, - dims=shape_value.shape, - vals=shape_value)) - reshape_node.input[1] = constant_shape_name - reshape_node.name = self.create_node_name('Reshape', 'Reshape_Fuse') - nodes_to_remove.extend([cast_node, concat_node]) - nodes_to_add.append(new_node) - - def __fuse_reshape_after_sotfmax(self, reshape_node, nodes_to_remove, nodes_to_add): - # Check that it is reshape after softmax. - path = self.match_parent_path(reshape_node, ['Transpose', 'MatMul', 'Softmax', 'Add'], [0, 0, 0, 0]) - if path is None: - return - - path0 = self.match_parent_path(reshape_node, ['Cast', 'Concat', 'Unsqueeze', 'Mul'], [1, 0, 0, 0]) - if path0 is None: - return - cast_node, concat_node, unsqueeze_node, mul_node = path0 - - # Verify that cast has attribute "to" = 7 - is_good_cast = False - for att in cast_node.attribute: - if att.name == 'to' and att.i == 7: - is_good_cast = True - break - if not is_good_cast: - return - - if not len(concat_node.input) == 2: - return - - shape = [0,0] - - if self.get_initializer(concat_node.input[1]): - concat_1 = self.get_initializer(concat_node.input[1]) - shape.extend(numpy_helper.to_array(concat_1)) - else: - return - - shape_value = np.asarray(shape, dtype=np.int64) - - constant_shape_name = self.create_node_name('Constant', 'constant_shape') - new_node = onnx.helper.make_node('Constant', - inputs=[], - outputs=[constant_shape_name], - value=onnx.helper.make_tensor(name='const_tensor', - data_type=TensorProto.INT64, - dims=shape_value.shape, - vals=shape_value)) - reshape_node.input[1] = constant_shape_name - reshape_node.name = self.create_node_name('Reshape', 'Reshape_Fuse') - nodes_to_remove.extend([cast_node, concat_node, unsqueeze_node, mul_node]) - nodes_to_add.append(new_node) - - def __fuse_reshape_after_normalize(self, reshape_node, nodes_to_remove): - parent = self.get_parent(reshape_node, 0) - if parent is None or not parent.op_type == self.normalize_name: - return - - parent_path = self.match_parent_path( - reshape_node, - ['Cast', 'Concat', 'Unsqueeze', 'Cast', 'Squeeze', 'Slice', 'Cast', 'Shape'], - [1, 0, 0, 0, 0, 0, 0, 0]) - - if not parent_path is None: - nodes_to_remove.extend(parent_path) - - nodes_to_remove.append(reshape_node) - - self.replace_input_of_all_nodes(reshape_node.output[0], reshape_node.input[0]) - - - - - def fuse_reshape(self): - nodes = self.nodes() - input_name_to_nodes = self.input_name_to_nodes() - output_name_to_node = self.output_name_to_node() - - nodes_to_remove = [] - nodes_to_add = [] - - for reshape_node in self.get_nodes_by_op_type('Reshape'): - self.__fuse_reshape_after_qkv(reshape_node, nodes_to_remove, nodes_to_add) - self.__fuse_reshape_after_sotfmax(reshape_node, nodes_to_remove, nodes_to_add) - self.__fuse_reshape_after_normalize(reshape_node, nodes_to_remove) - - logger.info(f"Count of nodes removed for Reshape fuse:{len(nodes_to_remove)}") - - self.remove_nodes(nodes_to_remove) - self.add_nodes(nodes_to_add) - """ Batch Layer Norm from Keras in Tensorflow: +----------------------+ @@ -333,7 +215,6 @@ class BertOnnxModelTF(BertOnnxModel): output_name_to_node = self.output_name_to_node() nodes_to_remove = [] - skip_layernorm_nodes = [] layernorm_nodes = [] for node in self.nodes(): if node.op_type == 'Add': @@ -388,29 +269,15 @@ class BertOnnxModelTF(BertOnnxModel): weight_input = mul_node_1.input[1] bias_input = sub_node_0.input[0] - if root_node.op_type == 'Add': - subgraph_nodes.append(root_node) - if not self.is_safe_to_fuse_nodes(subgraph_nodes, node.output, self.input_name_to_nodes(), self.output_name_to_node()): - subgraph_nodes.pop() - else: - nodes_to_remove.append(root_node) - normalize_node = onnx.helper.make_node(self.normalize_name, - inputs=[root_node.input[0], root_node.input[1], weight_input, bias_input], - outputs=[node.output[0]], - name=self.create_node_name(self.normalize_name, name_prefix="SkipLayerNorm")) - normalize_node.domain = "com.microsoft" - skip_layernorm_nodes.extend([normalize_node]) - continue - - normalize_node = onnx.helper.make_node('LayerNormalization', + normalize_node = onnx.helper.make_node( + 'LayerNormalization', inputs=[reduce_mean_node_1.input[0], weight_input, bias_input], - outputs=[node.output[0]], epsilon=epsilon) + outputs=[node.output[0]]) + normalize_node.attribute.extend([onnx.helper.make_attribute("epsilon", float(epsilon))]) layernorm_nodes.extend([normalize_node]) self.remove_nodes(nodes_to_remove) - self.add_nodes(skip_layernorm_nodes) self.add_nodes(layernorm_nodes) - logger.info(f"Fused SkipLayerNormalization count: {len(skip_layernorm_nodes)}") logger.info(f"Fused LayerNormalization count: {len(layernorm_nodes)}") def remove_identity(self): @@ -423,74 +290,6 @@ class BertOnnxModelTF(BertOnnxModel): self.remove_nodes(nodes_to_remove) logger.info(f"Removed Identity count: {len(nodes_to_remove)}") - def fuse_word_embedding(self): - nodes_to_remove = [] - for node in self.nodes(): - if node.op_type == 'Reshape': - data_path = self.match_parent_path(node, ['Gather', 'Reshape', 'Reshape'], [0, 1, 0]) - if data_path is None: - continue - gather_node, reshape_node_0, reshape_node_1 = data_path - shape_path = self.match_parent_path( - node, - ['Cast', 'Concat', 'Unsqueeze', 'Cast', 'Squeeze', 'Slice', 'Cast', 'Shape', 'Reshape'], - [ 1, 0, 0, 0, 0, 0, 0, 0, 0]) - - if shape_path is None: - continue - - cast_node_0, concat_node, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node, reshape_node_2 = shape_path - - if not reshape_node_1 == reshape_node_2: - continue - gather_node.input[1] = reshape_node_1.input[0] - self.replace_input_of_all_nodes(node.output[0], gather_node.output[0]) - nodes_to_remove.extend([reshape_node_0, cast_node_0, concat_node, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node, reshape_node_2, node]) - - self.remove_nodes(nodes_to_remove) - logger.info("Fused word embedding" if len(nodes_to_remove) > 0 else "Failed to fuse word embedding") - - def fuse_segment_embedding(self): - nodes_to_remove = [] - for node in self.nodes(): - if node.op_type == 'Reshape': - data_path = self.match_parent_path(node, ['MatMul', 'OneHot', 'Reshape'], [0, 0, 0]) - if data_path is None: - continue - - matmul_node, onehot_node, reshape_node_0 = data_path - - subgraph_nodes = [matmul_node, onehot_node, reshape_node_0] - if self.get_initializer(onehot_node.input[2]) is None: - concat_node_0 = self.get_parent(onehot_node, 2) - if concat_node_0 is None or concat_node_0.op_type != 'Concat': - continue - subgraph_nodes.append(concat_node_0) - - shape_path = self.match_parent_path( - node, - ['Cast', 'Concat', 'Unsqueeze', 'Cast', 'Squeeze', 'Slice', 'Cast', 'Shape', 'Gather'], - [ 1, 0, 0, 0, 0, 0, 0, 0, 0]) - - if shape_path is None: - continue - - cast_node_0, concat_node_1, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node, gather_node = shape_path - - gather_node = onnx.helper.make_node( - 'Gather', - inputs=[matmul_node.input[1], reshape_node_0.input[0]], - outputs=node.output, - name='segment_embedding_gather') - - nodes_to_remove.extend(subgraph_nodes) - - nodes_to_remove.extend([cast_node_0, concat_node_1, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node, node]) - self.add_node(gather_node) - - self.remove_nodes(nodes_to_remove) - logger.info("Fused segment embedding" if len(nodes_to_remove) > 0 else "Failed to fuse segment embedding") - def fuse_mask(self): nodes_to_remove = [] for node in self.nodes(): @@ -537,10 +336,252 @@ class BertOnnxModelTF(BertOnnxModel): self.add_node(unsqueeze_added_2) self.remove_nodes(nodes_to_remove) + if len(nodes_to_remove) > 0: + logger.info("Fused mask") + else: + self.fuse_mask_2() + + def fuse_mask_2(self): + nodes_to_remove = [] + for node in self.nodes(): + if node.op_type == 'Mul' and self.has_constant_input(node, -10000): + mask_path = self.match_parent_path( + node, + ['Sub', 'Unsqueeze', 'Mul', 'Cast', 'Reshape', 'Cast'], + [ 0, 1, 0, 1, 0, 0]) + if mask_path is None: + continue + sub_node, unsqueeze_node, mul_node, cast_node_0, reshape_node_0, cast_node_1 = mask_path + + mask_input_name = next(iter(self.mask_indice)) + if cast_node_1.input[0] != mask_input_name: + print("Cast input {} is not mask input{}".format(cast_node_1.input[0], mask_input_name)) + continue + + unsqueeze_added_1 = onnx.helper.make_node( + 'Unsqueeze', + inputs=[mask_input_name], + outputs=['mask_fuse_unsqueeze1_output'], + name='Mask_UnSqueeze_1', + axes=[1]) + + unsqueeze_added_2 = onnx.helper.make_node( + 'Unsqueeze', + inputs=['mask_fuse_unsqueeze1_output'], + outputs=['mask_fuse_unsqueeze2_output'], + name='Mask_UnSqueeze_2', + axes=[2]) + + cast_node_2 = onnx.helper.make_node( + 'Cast', + inputs=['mask_fuse_unsqueeze2_output'], + outputs=['mask_fuse_cast_output']) + cast_node_2.attribute.extend([onnx.helper.make_attribute("to", 1)]) + self.replace_node_input(sub_node, sub_node.input[1], 'mask_fuse_cast_output') + + nodes_to_remove.extend([unsqueeze_node, mul_node, cast_node_0, reshape_node_0, cast_node_1]) + self.add_node(unsqueeze_added_1) + self.add_node(unsqueeze_added_2) + self.add_node(cast_node_2) + + self.remove_nodes(nodes_to_remove) + + # Prune graph is done after removing nodes to remove island nodes. + if len(nodes_to_remove) > 0: + self.prune_graph() + logger.info("Fused mask" if len(nodes_to_remove) > 0 else "Failed to fuse mask") + def get_2d_initializers_from_parent_subgraphs(self, current_node): + """ + Find initializers that is 2D. Returns a dictionary with name as key and shape as value. + """ + parent_nodes = self.get_parent_subgraph_nodes(current_node, []) + initializers = {} + for node in parent_nodes: + for input in node.input: + initializer = self.get_initializer(input) + if initializer: + temp = numpy_helper.to_array(initializer) + if len(temp.shape) == 2: + initializers[initializer.name] = temp.shape + + return initializers + + def find_segment_ids(self, segment_embedding): + input_name_to_nodes = self.input_name_to_nodes() + if segment_embedding not in input_name_to_nodes: + return None + + nodes = input_name_to_nodes[segment_embedding] + if len(nodes) != 1: + return None + + graph_inputs = self.get_graph_inputs(nodes[0], recursive=True) + if len(graph_inputs) == 1: + return graph_inputs[0] + + print("Found multiple candidates of segment_ids", graph_inputs) + return None + + def find_input_ids(self, word_embedding): + input_name_to_nodes = self.input_name_to_nodes() + if word_embedding not in input_name_to_nodes: + return None + + nodes = input_name_to_nodes[word_embedding] + if len(nodes) != 1: + return None + + graph_inputs = self.get_graph_inputs(nodes[0], recursive=True) + if len(graph_inputs) == 1: + return graph_inputs[0] + + print("Found multiple candidates of input_ids", graph_inputs) + return None + + def find_mask_input(self, excluded_graph_inputs): + for node in self.nodes(): + if node.op_type == 'Softmax': + mask_path = self.match_parent_path( + node, + ['Add', 'Mul', 'Sub'], + [ 0, 1, None]) + if mask_path is None: + continue + add_node, mul_node, sub_node = mask_path + if self.has_constant_input(mul_node, -10000) and self.has_constant_input(sub_node, 1): + graph_inputs = self.get_graph_inputs(sub_node, recursive=True) + inputs = [input for input in graph_inputs if input not in excluded_graph_inputs] + if len(inputs) == 1: + return inputs[0] + + return None + + def create_embedding_subgraph(self, normalize_node, word_embedding, segment_embedding, position_embedding): + segment_ids = self.find_segment_ids(segment_embedding) + if segment_ids is None: + logger.info("Failed to find segment_ids. Cannot fuse embedding layer.") + return False + + input_ids = self.find_input_ids(word_embedding) + if input_ids is None: + logger.info("Failed to find input_ids. Cannot fuse embedding layer.") + return False + + mask_input = self.find_mask_input([segment_ids, input_ids]) + if mask_input is None: + logger.info("Failed to find input_mask. Cannot fuse embedding layer.") + return False + + self.bert_inputs = [input_ids, segment_ids, mask_input] + + mask_index = self.create_node_name('mask_index') + self.mask_indice[mask_input] = mask_index + + if self.find_graph_input(input_ids).type.tensor_type.elem_type != TensorProto.INT32: + casted, input_ids = self.cast_graph_input_to_int32(input_ids) + + if self.find_graph_input(segment_ids).type.tensor_type.elem_type != TensorProto.INT32: + casted, segment_ids = self.cast_graph_input_to_int32(segment_ids) + + if self.find_graph_input(mask_input).type.tensor_type.elem_type != TensorProto.INT32: + casted, mask_input = self.cast_graph_input_to_int32(mask_input) + + embed_output = self.create_node_name('embed_output') + embed_node = onnx.helper.make_node('EmbedLayerNormalization', + inputs=[input_ids, + segment_ids, + word_embedding, + position_embedding, + segment_embedding, + normalize_node.input[1], # gamma + normalize_node.input[2], # beta + mask_input], + outputs=[embed_output, mask_index], + name="EmbedLayer") + embed_node.domain = "com.microsoft" + self.replace_input_of_all_nodes(normalize_node.output[0], embed_output) + self.add_node(embed_node) + + def process_embedding(self): + """ + Automatically detect word, segment and position embeddings. + """ + logger.info("start processing embedding layer...") + output_name_to_node = self.output_name_to_node() + + layer_norm_nodes = self.get_nodes_by_op_type("LayerNormalization") + for layer_norm_node in layer_norm_nodes: + pos_embed_path = self.match_parent_path( + layer_norm_node, + ['Add', 'Reshape', 'Slice'], + [ 0, 1, 0], + output_name_to_node) + if pos_embed_path is None: + continue + + add_node, reshape_node, slice_node = pos_embed_path + initializer = self.get_initializer(slice_node.input[0]) + if initializer is None: + continue + + temp = numpy_helper.to_array(initializer) + if len(temp.shape) == 2: + logger.info("Found position embedding. name:{}, shape:{}".format(initializer.name, temp.shape)) + position_embedding = initializer.name + else: + logger.info("Failed to find position embedding. name:{}, shape:{}".format(initializer.name, temp.shape)) + return + + first_parent = self.get_parent(add_node, 0, output_name_to_node) + if first_parent is not None and first_parent.op_type == "Add": + embeddings = self.get_2d_initializers_from_parent_subgraphs(first_parent) + if len(embeddings) != 2: + logger.warning("Failed to find two embeddings (word and segment) from Add node. Found {}".format(embeddings)) + return + + word_embedding = None + segment_embedding = None + for name, shape in embeddings.items(): + if shape[0] == 2: + segment_embedding = name + logger.info("Found segment embedding. name:{}, shape:{}".format(name, shape)) + else: + word_embedding = name + logger.info("Found words embedding. name:{}, shape:{}".format(name, shape)) + + if word_embedding is None or segment_embedding is None: + logger.info("Failed to find both word and segment embedding") + return + + logger.info("Create Embedding node") + self.create_embedding_subgraph(layer_norm_node, word_embedding, segment_embedding, position_embedding) + # Prune graph to remove those original embedding nodes. + self.prune_graph() + break + def preprocess(self): self.remove_identity() - self.fuse_word_embedding() - self.fuse_segment_embedding() + self.process_embedding() + #TODO: remove fuse mask since we have embedding fused so fuse_attention shall handle the mask nodes. self.fuse_mask() + + def remove_reshape_before_first_attention(self): + attention_nodes = self.get_nodes_by_op_type("Attention") + for attention_node in attention_nodes: + path = self.match_parent_path( + attention_node, + ['Reshape', 'EmbedLayerNormalization'], + [ 0, 0]) + if path is None: + continue + logger.info("Remove Reshape before first Attention node.") + reshape, embed = path + self.replace_input_of_all_nodes(reshape.output[0], reshape.input[0]) + self.remove_node(reshape) + break + + def postprocess(self): + self.remove_reshape_before_first_attention() + self.prune_graph() diff --git a/onnxruntime/python/tools/bert/OnnxModel.py b/onnxruntime/python/tools/bert/OnnxModel.py index ef71fadbf7..c48c802bd1 100644 --- a/onnxruntime/python/tools/bert/OnnxModel.py +++ b/onnxruntime/python/tools/bert/OnnxModel.py @@ -4,20 +4,19 @@ #-------------------------------------------------------------------------- import logging -import onnx import sys import argparse import numpy as np from collections import deque +import onnx from onnx import ModelProto, TensorProto, numpy_helper logger = logging.getLogger(__name__) class OnnxModel: - def __init__(self, model, verbose): + def __init__(self, model): self.model = model self.node_name_counter = {} - self.verbose = verbose def input_name_to_nodes(self): input_name_to_nodes = {} @@ -84,7 +83,7 @@ class OnnxModel: for node in self.model.graph.node: OnnxModel.replace_node_output(node, old_output_name, new_output_name) - def get_initializer(self,name): + def get_initializer(self, name): for tensor in self.model.graph.initializer: if tensor.name == name: return tensor @@ -146,7 +145,8 @@ class OnnxModel: parent = output_name_to_node[input] if parent.op_type == parent_op_type and parent not in exclude: return parent, i - + else: + logger.debug(f"To find first {parent_op_type}, current {parent.op_type}") return None, None def match_parent(self, node, parent_op_type, input_index=None, output_name_to_node=None, exclude=[], return_indice=None): @@ -178,12 +178,16 @@ class OnnxModel: return parent if input_index >= len(node.input): + logger.debug(f"input_index {input_index} >= node inputs {len(node.input)}") return None parent = self.get_parent(node, input_index, output_name_to_node) if parent is not None and parent.op_type == parent_op_type and parent not in exclude: return parent + if parent is not None: + logger.debug(f"Expect {parent_op_type}, Got {parent.op_type}") + return None def match_parent_path(self, node, parent_op_types, parent_input_index, output_name_to_node=None, return_indice=None): @@ -211,6 +215,7 @@ class OnnxModel: for i, op_type in enumerate(parent_op_types): matched_parent = self.match_parent(current_node, op_type, parent_input_index[i], output_name_to_node, 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}") return None matched_parents.append(matched_parent) @@ -369,7 +374,6 @@ class OnnxModel: return full_name - def find_graph_input(self, input_name): for input in self.model.graph.input: if input.name == input_name: @@ -404,6 +408,23 @@ class OnnxModel: return unique_nodes + def get_graph_inputs(self, current_node, recursive=False): + """ + Find graph inputs that linked to current node. + """ + graph_inputs = [] + for input in current_node.input: + if self.find_graph_input(input) and input not in graph_inputs: + graph_inputs.append(input) + + if recursive: + parent_nodes = self.get_parent_subgraph_nodes(current_node, []) + for node in parent_nodes: + for input in node.input: + if self.find_graph_input(input) and input not in graph_inputs: + graph_inputs.append(input) + return graph_inputs + @staticmethod def input_index(node_output, child_node): index = 0 @@ -428,7 +449,57 @@ class OnnxModel: if len(unused_nodes) > 0: logger.info(f"Removed unused constant nodes: {len(unused_nodes)}") - def update_graph(self): + def prune_graph(self, outputs=None): + """ + Prune graph to keep only required outputs. It removes unnecessary inputs and nodes. + Nodes are not linked (directly or indirectly) to any required output will be removed. + + Args: + outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept. + """ + if outputs is None: + outputs = [output.name for output in self.model.graph.output] + + output_name_to_node = self.output_name_to_node() + all_nodes = [] + for output in outputs: + if output in output_name_to_node: + last_node = output_name_to_node[output] + if last_node in all_nodes: + continue + nodes = self.get_parent_subgraph_nodes(last_node, []) + all_nodes.append(last_node) + all_nodes.extend(nodes) + + nodes_to_remove = [] + for node in self.model.graph.node: + if node not in all_nodes: + nodes_to_remove.append(node) + + self.remove_nodes(nodes_to_remove) + + # remove outputs not in list + output_to_remove=[] + for output in self.model.graph.output: + if output.name not in outputs: + output_to_remove.append(output) + for output in output_to_remove: + self.model.graph.output.remove(output) + + # remove inputs not used by any node. + input_name_to_nodes = self.input_name_to_nodes() + input_to_remove=[] + for input in self.model.graph.input: + if input.name not in input_name_to_nodes: + input_to_remove.append(input) + for input in input_to_remove: + self.model.graph.input.remove(input) + + logger.info("Graph pruned: {} inputs, {} outputs and {} nodes are removed".format(len(input_to_remove), len(output_to_remove), len(nodes_to_remove))) + + self.update_graph() + + def update_graph(self, verbose=False): graph = self.model.graph remaining_input_names = [] @@ -437,8 +508,8 @@ class OnnxModel: for input_name in node.input: if input_name not in remaining_input_names: remaining_input_names.append(input_name) - if self.verbose: - logger.info(f"remaining input names: {remaining_input_names}" ) + if verbose: + logger.debug(f"remaining input names: {remaining_input_names}" ) # remove graph input that is not used inputs_to_remove = [] @@ -447,9 +518,9 @@ class OnnxModel: inputs_to_remove.append(input) for input in inputs_to_remove: graph.input.remove(input) - if self.verbose: - names_to_remove = [input.name for input in inputs_to_remove] - logger.info(f"remove {len(inputs_to_remove)} unused inputs: {names_to_remove}") + + names_to_remove = [input.name for input in inputs_to_remove] + logger.debug(f"remove {len(inputs_to_remove)} unused inputs: {names_to_remove}") # remove weights that are not used weights_to_remove = [] @@ -462,10 +533,10 @@ class OnnxModel: for initializer in weights_to_remove: graph.initializer.remove(initializer) - if self.verbose: - names_to_remove = [initializer.name for initializer in weights_to_remove] - logger.info(f"remove {len(weights_to_remove)} unused initializers: {names_to_remove}") - logger.info(f"remaining initializers:{weights_to_keep}") + names_to_remove = [initializer.name for initializer in weights_to_remove] + logger.debug(f"remove {len(weights_to_remove)} unused initializers: {names_to_remove}") + if verbose: + logger.debug(f"remaining initializers:{weights_to_keep}") self.remove_unused_constant() @@ -478,7 +549,17 @@ class OnnxModel: if output_to_remove in input_name_to_nodes: for impacted_node in input_name_to_nodes[output_to_remove]: if impacted_node not in nodes_to_remove: - if self.verbose: - logger.warning(f"it is not safe to remove nodes since output {output_to_remove} is used by {impacted_node}") + logger.debug(f"it is not safe to remove nodes since output {output_to_remove} is used by {impacted_node}") return False - return True \ No newline at end of file + return True + + def save_model_to_file(self, output_path): + logger.info(f"Output model to {output_path}") + + if output_path.endswith(".json"): + assert isinstance(self.model, ModelProto) + with open(output_path, "w") as out: + out.write(str(self.model)) + else: + with open(output_path, "wb") as out: + out.write(self.model.SerializeToString()) diff --git a/onnxruntime/python/tools/bert/README.md b/onnxruntime/python/tools/bert/README.md index 70e279820d..d3e8162a4a 100644 --- a/onnxruntime/python/tools/bert/README.md +++ b/onnxruntime/python/tools/bert/README.md @@ -33,7 +33,13 @@ def export_onnx(args, model, output_path): ``` ## Convert an BERT model from Tensorflow -Please refer to this notebook: https://github.com/onnx/tensorflow-onnx/blob/master/tutorials/BertTutorial.ipynb + +The tf2onnx and keras2onnx tools can be used to convert model that trained by Tensorflow. + +For Keras2onnx, please refere to its [example script](https://github.com/onnx/keras-onnx/blob/master/applications/nightly_build/test_transformers.py). + +For tf2onnx, please refer to this notebook: https://github.com/onnx/tensorflow-onnx/blob/master/tutorials/BertTutorial.ipynb + ## Model Optimization @@ -48,8 +54,8 @@ See below for description of all the options: - **input**: input model path - **output**: output model path -- **framework**: - Original framework. Only support TensorFlow and PyTorch +- **model_type**: (*defaul: bert*) + There are 3 model types: *bert*, *bert_tf* and *bert_keras* for models exported by PyTorch, tf2onnx and keras2onnx respectively. - **num_heads**: (*default: 12*) Number of attention heads, like 24 for BERT-large model. - **hidden_size**: (*default: 768*) @@ -63,5 +69,16 @@ See below for description of all the options: By default, model uses float32 in computation. If this flag is specified, half-precision float will be used. This option is recommended for NVidia GPU with Tensor Core like V100 and T4. For older GPUs, float32 is likely faster. - **verbose**: (*optional*) Print verbose information when this flag is specified. -- **run_onnxruntime**: (*optional*) - Use onnxruntime to do optimization. This option is only avaiable for pytorch model right now. + +## Supported Models + +Right now, this tool assumes input model has 3 inputs for input IDs, segment IDs, and attention mask. A model with less or addtional inputs might not be optimized. + +Most optimizations require exact match of a subgraph. That means this tool could only support similar models with such subgraphs. Any layout change in subgraph might cause optimization not working. Note that different training or export tool (including different versions) might get different graph layouts. + +Here is list of models that have been tested using this tool: +- **BertForSequenceClassification** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_glue.py) exported by PyTorch 1.2-1.4 using opset version 10 or 11. +- **BertForQuestionAnswering** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_squad.py) exported by PyTorch 1.2-1.4 using opset version 10 or 11. +- **TFBertForSequenceClassification** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_tf_glue.py) exported by keras2onnx 1.6.0. + +If your model is not in the list, the optimized model might not work. You are welcome to update the scripts to support new models. diff --git a/onnxruntime/python/tools/bert/bert_model_optimization.py b/onnxruntime/python/tools/bert/bert_model_optimization.py index 57fabf8434..bb8ff03e63 100644 --- a/onnxruntime/python/tools/bert/bert_model_optimization.py +++ b/onnxruntime/python/tools/bert/bert_model_optimization.py @@ -33,12 +33,32 @@ from onnx import ModelProto, TensorProto, numpy_helper import onnxruntime from BertOnnxModel import BertOnnxModel from BertOnnxModelTF import BertOnnxModelTF +from BertOnnxModelKeras import BertOnnxModelKeras logger = logging.getLogger('') -def run_onnxruntime(onnx_model_path, use_gpu, optimized_model_path=None): +# Map model type to tuple: optimizer class, export tools (pytorch, tf2onnx, keras2onnx) and whether OnnxRuntime has the optimization. +MODEL_CLASSES = { + "bert" : (BertOnnxModel, "pytorch", True), + "bert_tf": (BertOnnxModelTF, "tf2onnx", False), + "bert_keras" : (BertOnnxModelKeras, "keras2onnx", False) +} + +def optimize_by_onnxruntime(onnx_model_path, use_gpu, optimized_model_path=None): + """ + Use onnxruntime package to optimize model. It could support models exported by PyTorch. + + Args: + onnx_model_path (str): th path of input onnx model. + use_gpu (bool): whether the optimized model is targeted to run in GPU. + optimized_model_path (str or None): the path of optimized model. + + Returns: + optimized_model_path: the path of optimized model + """ if use_gpu and 'CUDAExecutionProvider' not in onnxruntime.get_available_providers(): logger.error("There is no gpu for onnxruntime to do optimization.") + return onnx_model_path sess_options = onnxruntime.SessionOptions() sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED @@ -61,85 +81,79 @@ def run_onnxruntime(onnx_model_path, use_gpu, optimized_model_path=None): def parse_arguments(): parser = argparse.ArgumentParser() - parser.add_argument('--input', required=True, type=str) - parser.add_argument('--output', required=True, type=str) - parser.add_argument('--framework', required=True, type=str.lower, choices=["tensorflow", "pytorch"], help="Original framework") + parser.add_argument('--input', required=True, type=str, + help="input onnx model path") - # model parameters. - parser.add_argument('--num_heads', required=False, type=int, default=12, help="number of attention heads") - parser.add_argument('--hidden_size', required=False, type=int, default=768) - parser.add_argument('--sequence_length', required=False, type=int, default=128) + parser.add_argument('--output', required=True, type=str, + help="optimized onnx model path") - # Use int32 (instead of int64) tensor as input to avoid unnecessary data - # type cast. - parser.add_argument('--input_int32', required=False, action='store_true') + parser.add_argument('--model_type', required=False, type=str.lower, default="bert", + choices=list(MODEL_CLASSES.keys()), + help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())) + + parser.add_argument('--num_heads', required=False, type=int, default=12, + help="number of attention heads") + + parser.add_argument('--hidden_size', required=False, type=int, default=768, + help="bert model hidden size. 768 for base model and 1024 for large") + + parser.add_argument('--sequence_length', required=False, type=int, default=128, + help="max sequence length") + + parser.add_argument('--input_int32', required=False, action='store_true', + help="Use int32 (instead of int64) tensor as input to avoid unnecessary data cast") parser.set_defaults(input_int32=False) - # For NVidia GPU with Tensor Core like V100 and T4, half-precision float - # brings better performance. - parser.add_argument('--float16', required=False, action='store_true') + parser.add_argument('--float16', required=False, action='store_true', + help="If your target device is V100 or T4 GPU, use this to convert float32 to float16 for best performance") parser.set_defaults(float16=False) - parser.add_argument('--gpu_only', required=False, action='store_true') + parser.add_argument('--gpu_only', required=False, action='store_true', + help="whether the target device is gpu or not") parser.set_defaults(gpu_only=False) parser.add_argument('--verbose', required=False, action='store_true') parser.set_defaults(verbose=False) - parser.add_argument('--use_onnxruntime', required=False, action='store_true') - parser.set_defaults(use_onnxruntime=False) - args = parser.parse_args() return args -def optimize_model(input, framework, gpu_only, num_heads, hidden_size, sequence_length, input_int32, float16, verbose): +def optimize_model(input, model_type, gpu_only, num_heads, hidden_size, sequence_length, input_int32, float16): + (optimizer_class, framework, run_onnxruntime) = MODEL_CLASSES[model_type] + + input_model_path = input + if run_onnxruntime: + input_model_path = optimize_by_onnxruntime(input_model_path, gpu_only) + logger.info("Use OnnxRuntime to optimize and save the optimized model to {}".format(input_model_path)) + model = ModelProto() - with open(input, "rb") as f: + with open(input_model_path, "rb") as f: model.ParseFromString(f.read()) - if framework == 'tensorflow': - bert_model = BertOnnxModelTF(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only, verbose) - else: #framework == 'pytorch' - bert_model = BertOnnxModel(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only, verbose) - + bert_model = optimizer_class(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only) bert_model.optimize() - - return bert_model -def output_model(model, output): - if output.endswith(".json"): # output to JSON. Only for test purpose. - if isinstance(model, ModelProto): - with open(output, "w") as out: - out.write(str(model)) - logger.info("Output JSON to {}.json".format(output)) - else: - with open(output, "wb") as out: - out.write(model.SerializeToString()) - logger.info("Output final model to {}".format(output)) + return bert_model def main(): args = parse_arguments() # output logging to stdout log_handler = logging.StreamHandler(sys.stdout) - log_handler.setFormatter(logging.Formatter('%(asctime)s %(message)s')) - log_handler.setLevel(logging.DEBUG) + 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) logger.addHandler(log_handler) - logger.setLevel(logging.DEBUG) - - input_model_path = args.input - if args.use_onnxruntime: - if framework == 'tensorflow': - logger.warning("onnxruntime does not have optimization for tensorflow model. Ignore the option --use_onnxruntime.") - else: - input_model_path = run_onnxruntime(input_model_path, args.gpu_only) - - bert_model = optimize_model(input_model_path, args.framework, args.gpu_only, args.num_heads, args.hidden_size, args.sequence_length, args.input_int32, args.float16, args.verbose) - - output_model(bert_model.model, args.output) + logger.setLevel(logging_level) + 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) + bert_model.save_model_to_file(args.output) if __name__ == "__main__": main() diff --git a/onnxruntime/python/tools/bert/test_bert_optimization.py b/onnxruntime/python/tools/bert/test_bert_optimization.py index 60daf944b0..3509d9a695 100644 --- a/onnxruntime/python/tools/bert/test_bert_optimization.py +++ b/onnxruntime/python/tools/bert/test_bert_optimization.py @@ -16,15 +16,16 @@ from onnx import helper, TensorProto, ModelProto from onnx.helper import make_node, make_tensor_value_info import numpy as np from onnx import numpy_helper -from bert_model_optimization import optimize_model, run_onnxruntime +from bert_model_optimization import optimize_model, optimize_by_onnxruntime from OnnxModel import OnnxModel +BERT_TEST_MODELS = { + "bert_pytorch_0": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_0.onnx', + "bert_pytorch_1": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_1.onnx', + "bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx' +} + class TestBertOptimization(unittest.TestCase): - def get_model(self, framework, index): - if framework == "pytorch": - return 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_{}.onnx'.format(index) - else: - return 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_{}.onnx'.format(index) def verify_node_count(self, bert_model, expected_node_count): for op_type, count in expected_node_count.items(): @@ -33,14 +34,14 @@ class TestBertOptimization(unittest.TestCase): self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count) def test_pytorch_model_0_cpu_onnxruntime(self): - input = self.get_model("pytorch", 0) + input = BERT_TEST_MODELS['bert_pytorch_0'] output = 'temp.onnx' - run_onnxruntime(input, use_gpu=False, optimized_model_path=output) + optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output) model = ModelProto() with open(output, "rb") as f: model.ParseFromString(f.read()) os.remove(output) - bert_model = OnnxModel(model, verbose=False) + bert_model = OnnxModel(model) expected_node_count = { 'EmbedLayerNormalization': 1, 'Attention': 12, @@ -56,14 +57,14 @@ class TestBertOptimization(unittest.TestCase): print("skip test_pytorch_model_0_gpu_onnxruntime since no gpu found") return - input = self.get_model("pytorch", 0) + input = BERT_TEST_MODELS['bert_pytorch_0'] output = 'temp.onnx' - run_onnxruntime(input, use_gpu=True, optimized_model_path=output) + optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output) model = ModelProto() with open(output, "rb") as f: model.ParseFromString(f.read()) os.remove(output) - bert_model = OnnxModel(model, verbose=False) + bert_model = OnnxModel(model) expected_node_count = { 'EmbedLayerNormalization': 1, 'Attention': 12, @@ -75,14 +76,14 @@ class TestBertOptimization(unittest.TestCase): self.verify_node_count(bert_model, expected_node_count) def test_pytorch_model_1_cpu_onnxruntime(self): - input = self.get_model("pytorch", 1) + input = BERT_TEST_MODELS['bert_pytorch_1'] output = 'temp.onnx' - run_onnxruntime(input, use_gpu=False, optimized_model_path=output) + optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output) model = ModelProto() with open(output, "rb") as f: model.ParseFromString(f.read()) os.remove(output) - bert_model = OnnxModel(model, verbose=False) + bert_model = OnnxModel(model) expected_node_count = { 'EmbedLayerNormalization': 1, 'Attention': 12, @@ -99,14 +100,14 @@ class TestBertOptimization(unittest.TestCase): print("skip test_pytorch_model_1_gpu_onnxruntime since no gpu found") return - input = self.get_model("pytorch", 1) + input = BERT_TEST_MODELS['bert_pytorch_1'] output = 'temp.onnx' - run_onnxruntime(input, use_gpu=True, optimized_model_path=output) + optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output) model = ModelProto() with open(output, "rb") as f: model.ParseFromString(f.read()) os.remove(output) - bert_model = OnnxModel(model, verbose=False) + bert_model = OnnxModel(model) expected_node_count = { 'EmbedLayerNormalization': 1, 'Attention': 12, @@ -119,18 +120,18 @@ class TestBertOptimization(unittest.TestCase): self.verify_node_count(bert_model, expected_node_count) def test_pytorch_model_0_cpu(self): - input = self.get_model("pytorch", 0) - bert_model = optimize_model(input, framework='pytorch', gpu_only=False, + input = BERT_TEST_MODELS['bert_pytorch_0'] + bert_model = optimize_model(input, 'bert', gpu_only=False, num_heads=2, hidden_size=8, sequence_length=10, - input_int32=False, float16=False, verbose=False) + input_int32=False, float16=False) expected_node_count = { 'EmbedLayerNormalization': 1, 'Attention': 12, 'SkipLayerNormalization': 24, - 'Gelu': 12, + 'Gelu': 0, 'FastGelu': 0, - 'BiasGelu': 0 + 'BiasGelu': 12 } self.verify_node_count(bert_model, expected_node_count) @@ -139,10 +140,10 @@ class TestBertOptimization(unittest.TestCase): print("skip test_pytorch_model_0_gpu since no gpu found") return - input = self.get_model("pytorch", 0) - bert_model = optimize_model(input, framework='pytorch', gpu_only=True, + input = BERT_TEST_MODELS['bert_pytorch_0'] + bert_model = optimize_model(input, 'bert', gpu_only=True, num_heads=2, hidden_size=8, sequence_length=10, - input_int32=False, float16=False, verbose=False) + input_int32=False, float16=False) expected_node_count = { 'EmbedLayerNormalization': 1, @@ -154,25 +155,18 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count) - def test_tensorflow_model_1_cpu(self): - input = self.get_model("tensorflow", 1) + def test_keras_model_1_cpu(self): + input = BERT_TEST_MODELS['bert_keras_0'] - # The model need constant folding. Use onnxruntime to do so for now. - temp = 'temp.onnx' - run_onnxruntime(input, use_gpu=False, optimized_model_path=temp) - - bert_model = optimize_model(temp, framework='tensorflow', gpu_only=False, + bert_model = optimize_model(input, 'bert_keras', gpu_only=False, num_heads=2, hidden_size=8, sequence_length=7, - input_int32=False, float16=False, verbose=False) - os.remove(temp) + input_int32=False, float16=False) - # Optimization for tensorflow model is still on-going. - # TODO: update this after code complete. expected_node_count = { - 'EmbedLayerNormalization': 0, - 'Attention': 0, + 'EmbedLayerNormalization': 1, + 'Attention': 12, 'LayerNormalization': 0, - 'SkipLayerNormalization': 25, + 'SkipLayerNormalization': 24, 'BiasGelu': 0, 'Gelu': 12, 'FastGelu': 0