saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support distill-bert fusion in transformers tool (#4631)

Browse source 
* checkin attention

* checkin embedlayer but cause invalid onnx model

* resolve comments

* fix comments

* check return values

* add version limit

* fix comments

* add warning
This commit is contained in:
Ye Wang 2020-07-31 17:57:54 -07:00 committed by GitHub
parent 8cf2c1c410
commit b1bfff34e0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 160 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
onnxruntime/python/tools/transformers/fusion_attention.py
View file

@ -178,19 +178,33 @@ class FusionAttention(Fusion):
return
(_, _, add_v, matmul_v) = v_nodes
is_distill = False;
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Div', 'MatMul'], [0, 0, 0, 0])
if qk_nodes is None:
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, 0, 0])
if qk_nodes is None:
logger.debug("fuse_attention: failed to match qk path")
return
(_, add_qk, _, matmul_qk) = qk_nodes
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'MatMul', 'Div'], [0, 0, 2, 0])
is_distill = True
if qk_nodes is None:
logger.debug("fuse_attention: failed to match qk path")
return
add_qk = None
matmul_qk = None
where_qk = None
if is_distill:
(_, where_qk, matmul_qk, _) = qk_nodes
else:
(_, add_qk, _, matmul_qk) = qk_nodes
q_nodes = self.model.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")
return
(_, _, add_q, matmul_q) = q_nodes
q_nodes = self.model.match_parent_path(matmul_qk, ['Div', 'Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, 0, 0])
if q_nodes is None:
logger.debug("fuse_attention: failed to match q path")
return
add_q = q_nodes[-2]
matmul_q = q_nodes[-1]
k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0])
if k_nodes is None:
@ -203,9 +217,15 @@ class FusionAttention(Fusion):
matmul_k = k_nodes[-1]
# Note that Cast might be removed by OnnxRuntime so we match two patterns here.
_, mask_nodes, _ = self.model.match_parent_paths(
add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node)
mask_nodes = None
if is_distill:
_, mask_nodes, _ = self.model.match_parent_paths(
where_qk, [(['Expand', 'Reshape', 'Equal'], [0, 0, 0]),
(['Cast', 'Expand', 'Reshape', 'Equal'], [0, 0, 0, 0])], output_name_to_node)
else :
_, mask_nodes, _ = self.model.match_parent_paths(
add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node)
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
@ -228,4 +248,4 @@ class FusionAttention(Fusion):
# Use prune graph to remove mask nodes since they are shared by all attention nodes.
#self.nodes_to_remove.extend(mask_nodes)
self.prune_graph = True
self.prune_graph = True

211
onnxruntime/python/tools/transformers/fusion_embedlayer.py
View file

@ -41,60 +41,12 @@ class FusionEmbedLayerNoMask(Fusion):
self.utils = FusionUtils(model)
self.attention = None
def fuse(self, node, input_name_to_nodes, output_name_to_node):
if self.model.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is None:
logger.debug("Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather")
return
self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False)
if self.attention is None:
# In case user disables attention fusion, check whether subgraph looks like Attention.
if node.output[0] not in input_name_to_nodes:
return
children = input_name_to_nodes[node.output[0]]
children_types = sorted([child.op_type for child in children])
if children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']:
logger.debug("No Attention like subgraph in children of SkipLayerNormalization")
return
# Assume the order of embeddings are word_embedding + position_embedding + segment_embedding
normalize_node = node
word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0])
if word_embedding_path is None:
logger.info("Word embedding path is not found. Embed layer cannot be fused.")
return
add_node, word_embedding_gather = word_embedding_path
input_ids = word_embedding_gather.input[1]
position_embedding_expand = None
position_embedding_shape = None
position_embedding_path = self.model.match_parent_path(normalize_node, ['Reshape', 'Slice'], [1, 0])
if position_embedding_path is not None:
_, position_embedding_weight_node = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand', 'Shape'], [1, 1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, position_embedding_shape = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
add_node, ['Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [1, 1, 1, 1, 0, 0])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, _, _, _, position_embedding_shape = position_embedding_path
else:
# Here we will not try to get exact match. Instead, we only try identify position embedding weights.
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand'], [1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand = position_embedding_path
else:
logger.info("Position embedding path is not found. Embed layer cannot be fused.")
return
if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids:
logger.info("position and word embedding is expected to be applied on same input")
return
def match_segment_path(self, normalize_node, input_name_to_nodes, output_name_to_node, input_ids_cast_node):
segment_ids = None
segment_embedding_gather = None
segment_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [1])
if segment_embedding_path is None:
segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 1])
if segment_embedding_path is None:
@ -106,6 +58,99 @@ class FusionEmbedLayerNoMask(Fusion):
segment_ids = segment_embedding_gather.input[1]
self.nodes_to_remove.extend(segment_embedding_path)
if self.model.find_graph_input(segment_ids):
casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids)
else:
segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(segment_ids)
# Cast might be removed by OnnxRuntime.
_, segment_id_path, _ = self.model.match_parent_paths(
segment_ids_cast_node,
[(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]),
(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node)
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"]))
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])))
segment_ids = "zeros_for_input_shape"
return segment_ids, segment_embedding_gather
def fuse(self, node, input_name_to_nodes, output_name_to_node):
is_distill = False;
if self.model.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is None and self.model.match_parent_path(node, ['Gather'], [0]) is None:
logger.debug("Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather or SkipLayerNormalization[0] <-- Gather")
return
self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False)
if self.attention is None:
# In case user disables attention fusion, check whether subgraph looks like Attention.
if node.output[0] not in input_name_to_nodes:
return
children = input_name_to_nodes[node.output[0]]
children_types = sorted([child.op_type for child in children])
if children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization'] and children_types != ['MatMul', 'MatMul', 'MatMul', 'Shape', 'Shape', 'SkipLayerNormalization']:
logger.debug("No Attention like subgraph in children of SkipLayerNormalization")
return
# Assume the order of embeddings are word_embedding + position_embedding + segment_embedding
normalize_node = node
add_node = None
word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0])
if word_embedding_path is not None:
add_node, word_embedding_gather = word_embedding_path
else:
word_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [0])
if word_embedding_path is not None:
word_embedding_gather = word_embedding_path[0]
is_distill = True;
else:
logger.info("Word embedding path is not found. Embed layer cannot be fused.")
return
input_ids = word_embedding_gather.input[1]
position_embedding_expand = None
position_embedding_shape = None
position_embedding_path = self.model.match_parent_path(normalize_node, ['Gather', 'Expand'], [1, 1]) # for distill-bert
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(normalize_node, ['Reshape', 'Slice'], [1, 0])
if position_embedding_path is not None:
_, position_embedding_weight_node = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand', 'Shape'], [1, 1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, position_embedding_shape = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
add_node, ['Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [1, 1, 1, 1, 0, 0])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, _, _, _, position_embedding_shape = position_embedding_path
else:
# Here we will not try to get exact match. Instead, we only try identify position embedding weights.
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand'], [1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand = position_embedding_path
else:
logger.info("Position embedding path is not found. Embed layer cannot be fused.")
return
if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids:
logger.info("position and word embedding is expected to be applied on same input")
return
if position_embedding_expand and position_embedding_shape:
input_parent = self.model.get_parent(position_embedding_shape, 0, output_name_to_node)
subgraph_nodes = self.model.get_parent_subgraph_nodes(position_embedding_expand,
@ -115,7 +160,6 @@ class FusionEmbedLayerNoMask(Fusion):
self.nodes_to_remove.extend(word_embedding_path)
self.nodes_to_remove.extend(position_embedding_path)
self.nodes_to_remove.extend(segment_embedding_path)
self.nodes_to_remove.extend([normalize_node])
@ -126,41 +170,46 @@ class FusionEmbedLayerNoMask(Fusion):
else:
input_ids, input_ids_cast_node = self.utils.cast_input_to_int32(input_ids)
if self.model.find_graph_input(segment_ids):
casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids)
else:
segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(segment_ids)
# Cast might be removed by OnnxRuntime.
_, segment_id_path, _ = self.model.match_parent_paths(
segment_ids_cast_node,
[(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]),
(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node)
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"]))
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])))
segment_ids = "zeros_for_input_shape"
node_name = self.model.create_node_name('EmbedLayerNormalization')
output_name = node_name + "_output"
embed_node = helper.make_node(
'EmbedLayerNormalization',
inputs=[
embed_node_inputs = None
if is_distill == False:
segment_path = self.match_segment_path(normalize_node, input_name_to_nodes, output_name_to_node, input_ids_cast_node)
if segment_path is None:
return
else:
from packaging.version import Version
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')
return
segment_ids, segment_embedding_gather = segment_path
embed_node_inputs=[
input_ids,
segment_ids,
word_embedding_gather.input[0],
position_embedding_weight_node.input[0],
segment_embedding_gather.input[0],
normalize_node.input[2],
normalize_node.input[3] # gamma and beta
]
else:
embed_node_inputs=[
input_ids,
segment_ids,
'',
word_embedding_gather.input[0],
position_embedding_weight_node.input[0],
segment_embedding_gather.input[0],
'',
normalize_node.input[2],
normalize_node.input[3] # gamma and beta
],
]
embed_node = helper.make_node(
'EmbedLayerNormalization',
embed_node_inputs,
outputs=[node_name + "_output", node_name + "_dummy_mask_index"],
name=node_name)