mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Support fusions inside subgraphs in optimizer tool (#7701)
* skip subgraph when updating model * intreim checkin * interim checkin 2 * support transformers optimizations in subgraph * change more files * fix comments typo
This commit is contained in:
Ye Wang
GitHub
parent
ad95b19322
commit
5e8086ad8e
15 changed files with 168 additions and 54 deletions
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
return graph_inputs
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:])
|
||||
|
|
|
|||
Loading…
Reference in a new issue