diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index e6f9ec9df2..82395f92d4 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -197,7 +197,7 @@ class FusionAttention(Fusion): # Sometimes weights and bias are stored in fp16 if q_weight.data_type == 10: weight.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(weight).astype(np.float16), weight.name)) - self.model.add_initializer(weight) + self.model.add_initializer(weight, self.this_graph_name) bias = helper.make_tensor(name=attention_node_name + '_qkv_bias', data_type=TensorProto.FLOAT, @@ -205,7 +205,7 @@ class FusionAttention(Fusion): vals=qkv_bias.flatten().tolist()) if q_bias.data_type == 10: bias.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(bias).astype(np.float16), bias.name)) - self.model.add_initializer(bias) + self.model.add_initializer(bias, self.this_graph_name) attention_inputs = [input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias'] if mask_index is not None: @@ -362,6 +362,7 @@ class FusionAttention(Fusion): return self.nodes_to_add.append(new_node) + self.node_name_to_graph_name[new_node.name] = self.this_graph_name if einsum_node is not None: unique_index = einsum_node.input[0] @@ -372,10 +373,10 @@ class FusionAttention(Fusion): vals=np.int64([0, 0, num_heads, int(hidden_size / num_heads)]).tobytes(), raw=True) - self.model.add_initializer(shape_tensor) + self.model.add_initializer(shape_tensor, self.this_graph_name) self.model.add_node( helper.make_node("Reshape", [attention_last_node.output[0], shape_tensor.name], [new_edge], - "reshape_modified_" + unique_index)) + "reshape_modified_" + unique_index), self.this_graph_name) einsum_node.input[0] = new_edge self.nodes_to_remove.extend([attention_last_node, transpose_qkv, matmul_qkv]) diff --git a/onnxruntime/python/tools/transformers/fusion_base.py b/onnxruntime/python/tools/transformers/fusion_base.py index 863588dc06..08b8ee2990 100644 --- a/onnxruntime/python/tools/transformers/fusion_base.py +++ b/onnxruntime/python/tools/transformers/fusion_base.py @@ -5,6 +5,7 @@ from logging import getLogger from onnx_model import OnnxModel from typing import Union, List +from onnx import GraphProto logger = getLogger(__name__) @@ -22,6 +23,8 @@ class Fusion: self.nodes_to_remove: List = [] self.nodes_to_add: List = [] self.prune_graph: bool = False + self.node_name_to_graph_name: dict = {} + self.this_graph_name: str = None def apply(self): logger.debug(f"start {self.description} fusion...") @@ -31,6 +34,10 @@ class Fusion: # This assumes that two search ops will not be fused at same time! for search_op_type in self.search_op_types: for node in self.model.get_nodes_by_op_type(search_op_type): + graph = self.model.get_graph_by_node(node) + if graph is None: + raise Exception("Can not find node in any graphs") + self.this_graph_name = graph.name self.fuse(node, input_name_to_nodes, output_name_to_node) op_list = [node.op_type for node in self.nodes_to_add] @@ -39,7 +46,7 @@ class Fusion: logger.info(f"Fused {self.description} count: {count}") self.model.remove_nodes(self.nodes_to_remove) - self.model.add_nodes(self.nodes_to_add) + self.model.add_nodes(self.nodes_to_add, self.node_name_to_graph_name) if self.prune_graph: self.model.prune_graph() diff --git a/onnxruntime/python/tools/transformers/fusion_biasgelu.py b/onnxruntime/python/tools/transformers/fusion_biasgelu.py index f66627c091..cfdf00b22f 100644 --- a/onnxruntime/python/tools/transformers/fusion_biasgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_biasgelu.py @@ -59,3 +59,4 @@ class FusionBiasGelu(Fusion): name=self.model.create_node_name(fuse_op_type, gelu_op_type + "_AddBias_")) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index cd9289d308..dbf403128f 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -37,7 +37,8 @@ class FusionEmbedLayerNoMask(Fusion): SkipLayerNormalization """ def __init__(self, model: OnnxModel, description='no mask'): - super().__init__(model, "EmbedLayerNormalization", ["SkipLayerNormalization", "LayerNormalization"], description) + super().__init__(model, "EmbedLayerNormalization", ["SkipLayerNormalization", "LayerNormalization"], + description) self.utils = FusionUtils(model) self.attention = None @@ -82,12 +83,14 @@ class FusionEmbedLayerNoMask(Fusion): 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.model.add_node( - helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"])) + helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"]), + self.this_graph_name) self.model.add_node( helper.make_node('ConstantOfShape', inputs=["input_shape"], outputs=["zeros_for_input_shape"], - value=helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1]))) + value=helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1])), + self.this_graph_name) segment_ids = "zeros_for_input_shape" return segment_ids, segment_embedding_gather @@ -100,7 +103,8 @@ class FusionEmbedLayerNoMask(Fusion): logger.debug( "Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather or SkipLayerNormalization[0] <-- Gather" ) - if node.op_type != "LayerNormalization" or self.model.match_parent_path(node, ['Add', 'Gather'], [0, 1]) is None: + if node.op_type != "LayerNormalization" or self.model.match_parent_path(node, ['Add', 'Gather'], + [0, 1]) is None: return self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) @@ -135,7 +139,8 @@ class FusionEmbedLayerNoMask(Fusion): import onnxruntime if Version(onnxruntime.__version__) <= Version("1.4.0"): logger.warning( - 'Please install onnxruntime with version > 1.4.0 for embedlayer fusion support for distilbert') + 'Please install onnxruntime with version > 1.4.0 for embedlayer fusion support for distilbert' + ) return else: logger.info("Word embedding path is not found. Embed layer cannot be fused.") @@ -175,7 +180,8 @@ class FusionEmbedLayerNoMask(Fusion): if position_embedding_path is not None: position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path else: - position_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather', 'Slice'], [0, 1, 1]) + position_embedding_path = self.model.match_parent_path( + normalize_node, ['Add', 'Gather', 'Slice'], [0, 1, 1]) if position_embedding_path is not None: _, position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path else: @@ -262,6 +268,7 @@ class FusionEmbedLayerNoMask(Fusion): self.model.replace_input_of_all_nodes(normalize_node.output[0], output_name) self.nodes_to_add.append(embed_node) + self.node_name_to_graph_name[embed_node.name] = self.this_graph_name class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): diff --git a/onnxruntime/python/tools/transformers/fusion_fastgelu.py b/onnxruntime/python/tools/transformers/fusion_fastgelu.py index a3f755fe74..c4685f4b45 100644 --- a/onnxruntime/python/tools/transformers/fusion_fastgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_fastgelu.py @@ -110,6 +110,7 @@ class FusionFastGelu(Fusion): name=self.model.create_node_name('FastGelu')) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True def fuse_2(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: @@ -204,9 +205,9 @@ class FusionFastGelu(Fusion): name=self.model.create_node_name('FastGelu')) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True - def fuse_3(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: """ OpenAI's gelu implementation, also used in Megatron: @@ -301,9 +302,10 @@ class FusionFastGelu(Fusion): self.nodes_to_remove.extend(subgraph_nodes) fused_node = helper.make_node('FastGelu', - inputs=[root_input], - outputs=mul_last.output, - name=self.model.create_node_name('FastGelu')) + inputs=[root_input], + outputs=mul_last.output, + name=self.model.create_node_name('FastGelu')) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True diff --git a/onnxruntime/python/tools/transformers/fusion_gelu.py b/onnxruntime/python/tools/transformers/fusion_gelu.py index 895ae8238e..0a7d657dd3 100644 --- a/onnxruntime/python/tools/transformers/fusion_gelu.py +++ b/onnxruntime/python/tools/transformers/fusion_gelu.py @@ -99,6 +99,7 @@ class FusionGelu(Fusion): fused_node = helper.make_node('Gelu', inputs=[subgraph_input], outputs=[subgraph_output]) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True def fuse_2(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: @@ -171,6 +172,7 @@ class FusionGelu(Fusion): fused_node = helper.make_node('Gelu', inputs=[root_node.output[0]], outputs=[mul.output[0]]) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True def fuse_3(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: @@ -237,4 +239,5 @@ class FusionGelu(Fusion): fused_node = helper.make_node('Gelu', inputs=[root_node.output[0]], outputs=[last_mul.output[0]]) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + self.node_name_to_graph_name[fused_node.name] = self.this_graph_name return True diff --git a/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py b/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py index 10cbac4d6c..253afabc6a 100644 --- a/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py +++ b/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py @@ -21,3 +21,4 @@ class FusionGeluApproximation(Fusion): new_node.domain = "com.microsoft" self.nodes_to_remove.append(node) self.nodes_to_add.append(new_node) + self.node_name_to_graph_name[new_node.name] = self.this_graph_name diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py index a9104b1b2f..cb5e01baea 100644 --- a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py @@ -46,6 +46,9 @@ class FusionGptAttention(Fusion): outputs=[output], name=attention_node_name + "_add") self.nodes_to_add.extend([attention_node, matmul_node, add_node]) + self.node_name_to_graph_name[attention_node.name] = self.this_graph_name + self.node_name_to_graph_name[matmul_node.name] = self.this_graph_name + self.node_name_to_graph_name[add_node.name] = self.this_graph_name def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): past = None diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py index 85aaf2680d..2c54cfbd81 100644 --- a/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py +++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py @@ -44,6 +44,9 @@ class FusionGptAttentionNoPast(Fusion): name=attention_node_name + "_add") self.nodes_to_add.extend([attention_node, matmul_node, add_node]) + self.node_name_to_graph_name[attention_node.name] = self.this_graph_name + self.node_name_to_graph_name[matmul_node.name] = self.this_graph_name + self.node_name_to_graph_name[add_node.name] = self.this_graph_name def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): return_indice = [] diff --git a/onnxruntime/python/tools/transformers/fusion_layernorm.py b/onnxruntime/python/tools/transformers/fusion_layernorm.py index d8f94948f5..0aa600aac8 100644 --- a/onnxruntime/python/tools/transformers/fusion_layernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_layernorm.py @@ -110,9 +110,12 @@ class FusionLayerNormalization(Fusion): normalize_node = helper.make_node('LayerNormalization', inputs=[node.input[0], weight_input, bias_input], - outputs=[last_add_node.output[0]]) + outputs=[last_add_node.output[0]], + name=self.model.create_node_name("LayerNormalization", + name_prefix="SkipLayerNorm")) normalize_node.attribute.extend([helper.make_attribute("epsilon", float(add_weight))]) self.nodes_to_add.append(normalize_node) + self.node_name_to_graph_name[normalize_node.name] = self.this_graph_name class FusionLayerNormalizationTF(Fusion): @@ -216,6 +219,8 @@ class FusionLayerNormalizationTF(Fusion): #TODO: add epsilon attribute fused_node = helper.make_node('LayerNormalization', inputs=[mul_node_3.input[0], weight_input, bias_input], - outputs=[node.output[0]]) + outputs=[node.output[0]], + name=self.model.create_node_name("LayerNormalization", + name_prefix="SkipLayerNorm")) fused_node.attribute.extend([helper.make_attribute("epsilon", float(epsilon))]) self.nodes_to_add.append(fused_node) diff --git a/onnxruntime/python/tools/transformers/fusion_reshape.py b/onnxruntime/python/tools/transformers/fusion_reshape.py index c190fc2be7..a325014767 100644 --- a/onnxruntime/python/tools/transformers/fusion_reshape.py +++ b/onnxruntime/python/tools/transformers/fusion_reshape.py @@ -136,3 +136,4 @@ class FusionReshape(Fusion): self.nodes_to_remove.extend(path2) self.nodes_to_remove.extend(path3) self.nodes_to_add.append(new_node) + self.node_name_to_graph_name[new_node.name] = self.this_graph_name diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 8c854cdca3..70a94da700 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -41,6 +41,8 @@ class FusionSkipLayerNormalization(Fusion): if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]): return else: + # shape_infer_helper can not handle subgraphs. Current work around is to disable skiplayernorm fusion + # longterm todo: support subgraph in symbolic_shape_infer or support add broadcasting in skiplayernorm op logger.warning( "symbolic shape infer failed. it's safe to ignore this message if there is no issue with optimized model" ) @@ -72,6 +74,7 @@ class FusionSkipLayerNormalization(Fusion): normalize_node.attribute.extend([helper.make_attribute("epsilon", 1.0E-12)]) self.nodes_to_add.append(normalize_node) + self.node_name_to_graph_name[normalize_node.name] = self.this_graph_name class FusionBiasSkipLayerNormalization(Fusion): @@ -136,3 +139,4 @@ class FusionBiasSkipLayerNormalization(Fusion): new_node.attribute.extend([helper.make_attribute("epsilon", 1.0E-12)]) self.nodes_to_add.append(new_node) + self.node_name_to_graph_name[new_node.name] = self.this_graph_name diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index f91f4951fe..9244421c4b 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -7,11 +7,10 @@ from typing import List, Tuple import logging import os import sys -import argparse from pathlib import Path import numpy as np from collections import deque -from onnx import ModelProto, TensorProto, numpy_helper, helper, external_data_helper, save_model +from onnx import onnx_pb, AttributeProto, ModelProto, TensorProto, numpy_helper, helper, external_data_helper, save_model from shape_infer_helper import SymbolicShapeInferenceHelper logger = logging.getLogger(__name__) @@ -22,8 +21,9 @@ class OnnxModel: self.model = model self.node_name_counter = {} self.shape_infer_helper = None + self.all_graphs = None - def infer_runtime_shape(self, dynamic_axis_mapping, update = False): + def infer_runtime_shape(self, dynamic_axis_mapping, update=False): shape_infer_helper = None if update: shape_infer_helper = SymbolicShapeInferenceHelper(self.model) @@ -36,13 +36,13 @@ class OnnxModel: if shape_infer_helper.infer(dynamic_axis_mapping): return shape_infer_helper except: - print("failed in shape inference", sys.exc_info()[0]) - + print("failed in shape inference", sys.exc_info()[0]) + return None def input_name_to_nodes(self): input_name_to_nodes = {} - for node in self.model.graph.node: + for node in self.nodes(): for input_name in node.input: if input_name not in input_name_to_nodes: input_name_to_nodes[input_name] = [node] @@ -52,36 +52,97 @@ class OnnxModel: def output_name_to_node(self): output_name_to_node = {} - for node in self.model.graph.node: + for node in self.nodes(): for output_name in node.output: output_name_to_node[output_name] = node return output_name_to_node def nodes(self): - return self.model.graph.node + all_nodes = [] + for graph in self.graphs(): + for node in graph.node: + all_nodes.append(node) + return all_nodes def graph(self): return self.model.graph + def graphs(self): + if self.all_graphs is not None: + return self.all_graphs + self.all_graphs = [] + graph_queue = [self.model.graph] + while graph_queue: + graph = graph_queue.pop(0) + self.all_graphs.append(graph) + for node in graph.node: + for attr in node.attribute: + if attr.type == AttributeProto.AttributeType.GRAPH: + assert (isinstance(attr.g, onnx_pb.GraphProto)) + graph_queue.append(attr.g) + if attr.type == AttributeProto.AttributeType.GRAPHS: + for g in attr.graphs: + assert (isinstance(g, onnx_pb.GraphProto)) + graph_queue.append(g) + return self.all_graphs + + def get_graph_by_node(self, node): + for graph in self.graphs(): + if node in graph.node: + return graph + return None + + def get_graph_by_name(self, graph_name): + for graph in self.graphs(): + if graph_name == graph.name: + return graph + return None + + def get_topological_insert_id(self, graph, outputs): + for idx, node in enumerate(graph.node): + for input in node.input: + if input in outputs: + return idx + return len(graph.node) + def remove_node(self, node): - if node in self.model.graph.node: - self.model.graph.node.remove(node) + for graph in self.graphs(): + if node in graph.node: + graph.node.remove(node) def remove_nodes(self, nodes_to_remove): for node in nodes_to_remove: self.remove_node(node) - def add_node(self, node): - self.model.graph.node.extend([node]) + def add_node(self, node, graph_name=None): + if graph_name is None or graph_name == self.model.graph.name: + self.model.graph.node.extend([node]) + else: + graph = self.get_graph_by_name(graph_name) + insert_idx = self.get_topological_insert_id(graph, node.output) + graph.node.insert(insert_idx, node) - def add_nodes(self, nodes_to_add): - self.model.graph.node.extend(nodes_to_add) + def add_nodes(self, nodes_to_add, node_name_to_graph_name=None): + if node_name_to_graph_name is None: + self.model.graph.node.extend(nodes_to_add) + else: + for node in nodes_to_add: + graph_name = node_name_to_graph_name[node.name] + self.add_node(node, graph_name) - def add_initializer(self, tensor): - self.model.graph.initializer.extend([tensor]) + def add_initializer(self, tensor, graph_name=None): + if graph_name is None or graph_name == self.model.graph.name: + self.model.graph.initializer.extend([tensor]) + else: + graph = self.get_graph_by_name(graph_name) + graph.initializer.extend([tensor]) - def add_input(self, input): - self.model.graph.input.extend([input]) + def add_input(self, input, graph_name=None): + if graph_name is None or graph_name == self.model.graph.name: + self.model.graph.input.extend([input]) + else: + graph = self.get_graph_by_name(graph_name) + graph.input.extend([input]) @staticmethod def replace_node_input(node, old_input_name, new_input_name): @@ -90,6 +151,7 @@ class OnnxModel: if node.input[j] == old_input_name: node.input[j] = new_input_name + # This function is deprecated since we use onnxconverter-common def replace_input_of_all_nodes(self, old_input_name, new_input_name): for node in self.model.graph.node: OnnxModel.replace_node_input(node, old_input_name, new_input_name) @@ -101,18 +163,24 @@ class OnnxModel: if node.output[j] == old_output_name: node.output[j] = new_output_name + # This function is deprecated since we use onnxconverter-common def replace_output_of_all_nodes(self, old_output_name, new_output_name): for node in self.model.graph.node: OnnxModel.replace_node_output(node, old_output_name, new_output_name) def get_initializer(self, name): - for tensor in self.model.graph.initializer: - if tensor.name == name: - return tensor + for graph in self.graphs(): + for tensor in graph.initializer: + if tensor.name == name: + return tensor return None def get_nodes_by_op_type(self, op_type): - return [n for n in self.model.graph.node if n.op_type == op_type] + nodes = [] + for node in self.nodes(): + if node.op_type == op_type: + nodes.append(node) + return nodes def get_children(self, node, input_name_to_nodes=None): if (input_name_to_nodes is None): @@ -644,6 +712,10 @@ class OnnxModel: remaining_input_names = [] for node in graph.node: + if node.op_type in ['Loop', 'Scan', 'If']: + # TODO: handle inner graph + logger.debug(f"Skip update_graph since graph has operator: {node.op_type}") + return if node.op_type != "Constant": for input_name in node.input: if input_name not in remaining_input_names: @@ -723,4 +795,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 \ No newline at end of file + return graph_inputs diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index 4fdfb1097e..d5c3f78264 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -49,7 +49,7 @@ class BertOptimizationOptions: class BertOnnxModel(OnnxModel): def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0): """Initialize BERT ONNX Model. - + Args: model (ModelProto): the ONNX model num_heads (int, optional): number of attentioin heads. Defaults to 0, and we will detect the parameter automatically. @@ -217,7 +217,6 @@ class BertOnnxModel(OnnxModel): def clean_graph(self): output_name_to_node = self.output_name_to_node() - nodes_to_add = [] nodes_to_remove = [] for node in self.nodes(): # Before: @@ -255,10 +254,9 @@ class BertOnnxModel(OnnxModel): name=node.name + "_remove_mask") attention_node.domain = "com.microsoft" attention_node.attribute.extend([helper.make_attribute("num_heads", self.num_heads)]) - nodes_to_add.append(attention_node) + self.add_node(attention_node, self.get_graph_by_node(attention_node).name) nodes_to_remove.append(node) self.remove_nodes(nodes_to_remove) - self.add_nodes(nodes_to_add) def postprocess(self): self.clean_graph() diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py b/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py index 9a682f833a..51522ba252 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert_tf.py @@ -87,14 +87,15 @@ class BertOnnxModelTF(BertOnnxModel): if segment_id_path and input_ids and input_ids == segment_id_path[-1].input[0]: logger.debug("Simplify semgent id path...") - self.add_node(helper.make_node('Shape', inputs=[input_ids], outputs=["input_shape"])) constantofshape_node = segment_id_path[0] + graph_name = self.get_graph_by_node(constantofshape_node).name + self.add_node(helper.make_node('Shape', inputs=[input_ids], outputs=["input_shape"]), graph_name) constantofshape_value = helper.get_attribute_value(constantofshape_node.attribute[0]) self.add_node( helper.make_node('ConstantOfShape', inputs=["input_shape"], outputs=["zeros_for_input_shape"], - value=constantofshape_value)) + value=constantofshape_value), graph_name) segment_ids = "zeros_for_input_shape" return segment_ids @@ -142,13 +143,15 @@ class BertOnnxModelTF(BertOnnxModel): logger.debug("Simplify semgent id path...") constantofshape_node = path_to_be_simplified[0] constantofshape_value = helper.get_attribute_value(constantofshape_node.attribute[0]) + graph_name = self.get_graph_by_node(constantofshape_node).name self.add_node( - helper.make_node('Shape', inputs=[duplicated_inputs[0]], outputs=["input_shape_for_mask"])) + helper.make_node('Shape', inputs=[duplicated_inputs[0]], outputs=["input_shape_for_mask"]), + graph_name) self.add_node( helper.make_node('ConstantOfShape', inputs=["input_shape_for_mask"], outputs=[unsqueeze_node.input[0]], - value=constantofshape_value)) + value=constantofshape_value), graph_name) return unsqueeze_node.input[0] return None @@ -203,7 +206,7 @@ class BertOnnxModelTF(BertOnnxModel): name="EmbedLayer") embed_node.domain = "com.microsoft" self.replace_input_of_all_nodes(normalize_node.output[0], embed_output) - self.add_node(embed_node) + self.add_node(embed_node, self.get_graph_by_node(normalize_node).name) def process_embedding(self): """ @@ -285,6 +288,8 @@ class BertOnnxModelTF(BertOnnxModel): start_nodes.extend(skip_layer_norm_nodes) start_nodes.extend(layer_norm_nodes) + graph_name = self.get_graph_by_node(start_nodes[0]).name + for normalize_node in start_nodes: # SkipLayerNormalization has two inputs, and one of them is the root input for attention. if normalize_node.op_type == 'LayerNormalization': @@ -366,7 +371,8 @@ class BertOnnxModelTF(BertOnnxModel): if squeeze_node is None and len(mask_nodes) == 5 and self.find_graph_input(mask_nodes[-1].input[0]) is None: mask_input = mask_nodes[-1].input[1] self.add_node( - helper.make_node("Squeeze", [mask_input], [squeeze_output_name], squeeze_node_name, axes=[1])) + helper.make_node("Squeeze", [mask_input], [squeeze_output_name], squeeze_node_name, axes=[1]), + graph_name) mask_nodes[-1].input[0] = squeeze_output_name is_same_root = self.check_attention_input(matmul_q, matmul_k, matmul_v, parent, output_name_to_node) @@ -391,13 +397,13 @@ class BertOnnxModelTF(BertOnnxModel): [[0, 0, self.num_heads, int(self.hidden_size / self.num_heads)]]).tobytes(), raw=True) - self.add_initializer(tensor) + self.add_initializer(tensor, graph_name) reshape_ = helper.make_node("Reshape", inputs=[attention_node.output[0], qkv_nodes[1].name + "_newshape"], outputs=[qkv_nodes[1].name + "_reshape_output"], name=qkv_nodes[1].name + "_reshape") qkv_nodes[1].input[0] = qkv_nodes[1].name + "_reshape_output" - self.add_node(reshape_) + self.add_node(reshape_, graph_name) if parent.op_type == 'Reshape': # Temporary work around: we require the skiplayernorm and attention op be fed with 3-d input hidden_size = numpy_helper.to_array(self.get_initializer(parent.input[1]))[1] @@ -406,10 +412,10 @@ class BertOnnxModelTF(BertOnnxModel): dims=[3], vals=np.int64([[1, -1, hidden_size]]).tobytes(), raw=True) - self.add_initializer(tensor) + self.add_initializer(tensor, graph_name) parent.input[1] = parent.name + "_modified" - self.add_node(attention_node) + self.add_node(attention_node, graph_name) attention_count += 1 nodes_to_remove.extend(qkv_nodes[2:])