From 3c9ece4a1106a4e206fdf825c55cd0c9c0608586 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 3 May 2021 20:42:13 -0700 Subject: [PATCH] [transformers optimizer] catch symbolic shape inference exception and clean up (#7560) catch symbolic shape inference exception. no prune graph when there is inner graph (Loop/If/Scan) add an wrapper for numpy_helper.to_array so that we can debug onnx graph without external data remove fuse_mask that is not used any more in onnx_model_bert_tf.py --- .../tools/transformers/fusion_attention.py | 22 ++-- .../tools/transformers/fusion_biasgelu.py | 5 +- .../transformers/fusion_skiplayernorm.py | 5 +- .../python/tools/transformers/fusion_utils.py | 16 ++- .../python/tools/transformers/onnx_model.py | 19 ++- .../tools/transformers/onnx_model_bert_tf.py | 109 +----------------- 6 files changed, 51 insertions(+), 125 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 7b866fc6bc..657a4e717f 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -9,7 +9,7 @@ from typing import Tuple, Union from onnx import helper, numpy_helper, TensorProto, NodeProto from onnx_model import OnnxModel from fusion_base import Fusion -from fusion_utils import FusionUtils +from fusion_utils import FusionUtils, NumpyHelper logger = getLogger(__name__) @@ -107,7 +107,7 @@ class FusionAttention(Fusion): logger.debug(f"{reshape_q.input[1]} is not initializer.") return self.num_heads, self.hidden_size # Fall back to user specified value - q_shape_value = numpy_helper.to_array(q_shape) + q_shape_value = NumpyHelper.to_array(q_shape) if len(q_shape_value) != 4 or (q_shape_value[2] <= 0 or q_shape_value[3] <= 0): logger.debug(f"q_shape_value={q_shape_value}. Expected value are like [0, 0, num_heads, head_size].") return self.num_heads, self.hidden_size # Fall back to user specified value @@ -145,7 +145,7 @@ class FusionAttention(Fusion): Returns: Union[NodeProto, None]: the node created or None if failed. """ - assert num_heads > 0 or hidden_size > 0 or (hidden_size % num_heads) == 0 + assert num_heads > 0 and hidden_size > 0 and (hidden_size % num_heads) == 0 q_weight = self.model.get_initializer(q_matmul.input[1]) k_weight = self.model.get_initializer(k_matmul.input[1]) @@ -159,9 +159,9 @@ class FusionAttention(Fusion): return None if not (k_weight and v_weight and q_bias and k_bias): return None - qw = numpy_helper.to_array(q_weight) - kw = numpy_helper.to_array(k_weight) - vw = numpy_helper.to_array(v_weight) + qw = NumpyHelper.to_array(q_weight) + kw = NumpyHelper.to_array(k_weight) + vw = NumpyHelper.to_array(v_weight) # Check if all matrices have the same shape assert qw.shape == kw.shape == vw.shape @@ -173,9 +173,9 @@ class FusionAttention(Fusion): qkv_weight = np.stack((qw, kw, vw), axis=1) - qb = numpy_helper.to_array(q_bias) - kb = numpy_helper.to_array(k_bias) - vb = numpy_helper.to_array(v_bias) + qb = NumpyHelper.to_array(q_bias) + kb = NumpyHelper.to_array(k_bias) + vb = NumpyHelper.to_array(v_bias) # 1d bias shape: [outsize,]. 2d bias shape: [a, b] where a*b = out_size assert qb.shape == kb.shape == vb.shape @@ -196,7 +196,7 @@ class FusionAttention(Fusion): # Sometimes weights and bias are stored in fp16 if q_weight.data_type == 10: - weight.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(weight).astype(np.float16), weight.name)) + weight.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(weight).astype(np.float16), weight.name)) self.model.add_initializer(weight) bias = helper.make_tensor(name=attention_node_name + '_qkv_bias', @@ -204,7 +204,7 @@ class FusionAttention(Fusion): dims=[3 * out_size], vals=qkv_bias.flatten().tolist()) if q_bias.data_type == 10: - bias.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(bias).astype(np.float16), bias.name)) + bias.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(bias).astype(np.float16), bias.name)) self.model.add_initializer(bias) attention_inputs = [input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias'] diff --git a/onnxruntime/python/tools/transformers/fusion_biasgelu.py b/onnxruntime/python/tools/transformers/fusion_biasgelu.py index d5d84058bb..f66627c091 100644 --- a/onnxruntime/python/tools/transformers/fusion_biasgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_biasgelu.py @@ -4,9 +4,10 @@ #-------------------------------------------------------------------------- from logging import getLogger -from onnx import helper, numpy_helper +from onnx import helper from onnx_model import OnnxModel from fusion_base import Fusion +from fusion_utils import NumpyHelper logger = getLogger(__name__) @@ -38,7 +39,7 @@ class FusionBiasGelu(Fusion): if initializer is None: continue bias_index = i - bias_weight = numpy_helper.to_array(initializer) + bias_weight = NumpyHelper.to_array(initializer) break if bias_weight is None: return diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 7292dc7b3d..8c854cdca3 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -4,9 +4,10 @@ #-------------------------------------------------------------------------- from logging import getLogger -from onnx import helper, numpy_helper +from onnx import helper from onnx_model import OnnxModel from fusion_base import Fusion +from fusion_utils import NumpyHelper logger = getLogger(__name__) @@ -99,7 +100,7 @@ class FusionBiasSkipLayerNormalization(Fusion): if initializer is None: continue bias_index = i - bias_weight = numpy_helper.to_array(initializer) + bias_weight = NumpyHelper.to_array(initializer) break if bias_weight is None: logger.debug(f"Bias weight not found") diff --git a/onnxruntime/python/tools/transformers/fusion_utils.py b/onnxruntime/python/tools/transformers/fusion_utils.py index 9d2ad5b902..ba352c120c 100644 --- a/onnxruntime/python/tools/transformers/fusion_utils.py +++ b/onnxruntime/python/tools/transformers/fusion_utils.py @@ -3,9 +3,10 @@ # Licensed under the MIT License. #-------------------------------------------------------------------------- from logging import getLogger -from onnx_model import OnnxModel from typing import Tuple -from onnx import helper, TensorProto +from onnx import helper, numpy_helper, TensorProto +from numpy import ndarray +from onnx_model import OnnxModel logger = getLogger(__name__) @@ -55,3 +56,14 @@ class FusionUtils: output_name = node.output[0] self.model.remove_node(node) self.model.replace_input_of_all_nodes(output_name, input_name) + +class NumpyHelper: + @staticmethod + def to_array(tensor:TensorProto, fill_zeros:bool = False) -> ndarray: + # When weights are in external data format but not presented, we can still test the optimizer with two changes: + # (1) set fill_zeros = True (2) change load_external_data=False in optimizer.py + if fill_zeros: + from onnx import mapping + return ndarray(shape=tensor.dims, dtype=mapping.TENSOR_TYPE_TO_NP_TYPE[tensor.data_type]) + + return numpy_helper.to_array(tensor) \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index ef8626a70e..f91f4951fe 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -32,9 +32,12 @@ class OnnxModel: if self.shape_infer_helper is None: self.shape_infer_helper = SymbolicShapeInferenceHelper(self.model) shape_infer_helper = self.shape_infer_helper - - if shape_infer_helper.infer(dynamic_axis_mapping): - return shape_infer_helper + try: + if shape_infer_helper.infer(dynamic_axis_mapping): + return shape_infer_helper + except: + print("failed in shape inference", sys.exc_info()[0]) + return None def input_name_to_nodes(self): @@ -585,6 +588,14 @@ class OnnxModel: Args: outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept. """ + + for node in self.model.graph.node: + # Some operators with inner graph in attributes like 'body' 'else_branch' or 'then_branch' + if node.op_type in ['Loop', 'Scan', 'If']: + # TODO: handle inner graph + logger.debug(f"Skip prune_graph since graph has operator: {node.op_type}") + return + if outputs is None: outputs = [output.name for output in self.model.graph.output] @@ -712,4 +723,4 @@ class OnnxModel: for input in self.model.graph.input: if self.get_initializer(input.name) is None: graph_inputs.append(input) - return graph_inputs + return graph_inputs \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py b/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py index 2d63dc7c0a..9a682f833a 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py @@ -42,106 +42,9 @@ class BertOnnxModelTF(BertOnnxModel): mask_nodes = self.match_parent_path(add_or_sub_before_softmax, ['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, None, 1, 0, 0]) + return mask_nodes - def fuse_mask(self): - nodes_to_remove = [] - for node in self.nodes(): - if node.op_type == 'Sub': - parent_path_constant = self.match_parent_path( - node, - ['Reshape', 'Mul', 'ConstantOfShape', 'Cast', 'Concat', 'Unsqueeze', 'Cast', 'Squeeze', 'Slice', 'Cast', 'Shape'], - [ 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # yapf: disable - if parent_path_constant is None: - continue - reshape_node_0, mul_node_0, constantofshape_node, cast_node_0, concat_node_0, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node = parent_path_constant - - parent_path_mask = self.match_parent_path( - mul_node_0, - ['Cast', 'Reshape', 'Cast', 'Concat', 'Unsqueeze'], - [ 1, 0, 1, 0, 0]) # yapf: disable - - if parent_path_mask is None: - continue - - cast_node_3, reshape_node_1, cast_node_4, concat_node_1, unsqueeze_node_1 = parent_path_mask - - if not unsqueeze_node_1 == unsqueeze_node: - continue - - unsqueeze_added_1 = onnx.helper.make_node('Unsqueeze', - inputs=[reshape_node_1.input[0]], - 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=[cast_node_3.input[0]], - name='Mask_UnSqueeze_2', - axes=[2]) - node.input[1] = cast_node_3.output[0] - - nodes_to_remove.extend([ - reshape_node_0, mul_node_0, constantofshape_node, cast_node_0, concat_node_0, unsqueeze_node, - cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node - ]) - nodes_to_remove.extend([reshape_node_1, cast_node_4, concat_node_1]) - self.add_node(unsqueeze_added_1) - 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', 'Cast', 'Slice', 'Unsqueeze'], [0, 1, 0, 0]) - if mask_path is None: - continue - sub_node, cast_node, slice_node, unsqueeze_node = mask_path - - mask_input_name = self.attention_mask.get_first_mask() - if unsqueeze_node.input[0] != mask_input_name: - print("Cast input {} is not mask input {}".format(unsqueeze_node.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]) - - #self.replace_node_input(cast_node, cast_node.input[0], 'mask_fuse_unsqueeze2_output') - 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([slice_node, unsqueeze_node, cast_node]) - 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. @@ -432,6 +335,7 @@ class BertOnnxModelTF(BertOnnxModel): if q_nodes is None: logger.debug("Failed to match q path") continue + add_q = q_nodes[-2] matmul_q = q_nodes[-1] @@ -469,11 +373,12 @@ class BertOnnxModelTF(BertOnnxModel): if is_same_root: mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0]) logger.debug("Create an Attention node.") + # For tf models, q and v are flipped. attention_node = self.attention_fusion.create_attention_node(mask_index, matmul_k, matmul_q, matmul_v, add_k, add_q, add_v, self.num_heads, self.hidden_size, parent.output[0], - qkv_nodes[2].output[0]) + qkv_nodes[2].output[0]) if attention_node is None: continue @@ -504,7 +409,6 @@ class BertOnnxModelTF(BertOnnxModel): self.add_initializer(tensor) parent.input[1] = parent.name + "_modified" - self.add_node(attention_node) attention_count += 1 @@ -524,8 +428,6 @@ class BertOnnxModelTF(BertOnnxModel): def preprocess(self): self.remove_identity() self.process_embedding() - #TODO: remove fuse mask since we have embedding fused so fuse_attention shall handle the mask nodes. - # self.fuse_mask() self.skip_reshape() def skip_reshape(self): @@ -554,5 +456,4 @@ class BertOnnxModelTF(BertOnnxModel): def postprocess(self): self.remove_reshape_before_first_attention() - # Temporary work around for the following comment as it will cause topological issues for a bert model - # self.prune_graph() + self.prune_graph()