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

Bert optimization for onnx model exported by keras2onnx (#3014)

Browse source 
* Optimization for Bert and DistilBert model exported by keras2onnx
* Add model_type parameter for models from different export tools (pytorch, tf2onnx, keras2onnx).
* Split LayerNormalization and SkipLayerNormalization fusions
This commit is contained in:
Tianlei Wu 2020-02-15 23:59:49 -08:00 committed by GitHub
parent 3626c46fad
commit aea76b0786
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 843 additions and 387 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

170
onnxruntime/python/tools/bert/BertOnnxModel.py
View file

@ -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}")

265
onnxruntime/python/tools/bert/BertOnnxModelKeras.py Normal file
View file

@ -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()

459
onnxruntime/python/tools/bert/BertOnnxModelTF.py
View file

@ -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()

119
onnxruntime/python/tools/bert/OnnxModel.py
View file

@ -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
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())

27
onnxruntime/python/tools/bert/README.md
View file

@ -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.

116
onnxruntime/python/tools/bert/bert_model_optimization.py
View file

@ -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()

74
onnxruntime/python/tools/bert/test_bert_optimization.py
View file

@ -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