mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Update bert optimization script for SQuAD model exported by keras2onnx (#3163)
Update script to make it work on fine-tuned bert model exported by keras2onnx
This commit is contained in:
Tianlei Wu
GitHub
parent
876d0c5430
commit
51a8c82908
13 changed files with 571 additions and 53 deletions
|
|
@ -685,7 +685,8 @@ class BertOnnxModel(OnnxModel):
|
|||
break
|
||||
|
||||
if normalize_node is None:
|
||||
logger.info("Failed to find embedding layer")
|
||||
if len(self.get_nodes_by_op_type("EmbedLayerNormalization")) == 0:
|
||||
logger.info("Failed to find embedding layer")
|
||||
return
|
||||
|
||||
# Here we assume the order of embedding is word_embedding +
|
||||
|
|
|
|||
|
|
@ -18,6 +18,41 @@ 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 match_mask_path(self, add_or_sub_before_softmax):
|
||||
mask_nodes = self.match_parent_path(add_or_sub_before_softmax,
|
||||
['Mul', 'Sub', 'Reshape', 'Cast'],
|
||||
[ 1, None, 1, 0])
|
||||
if mask_nodes is not None:
|
||||
return mask_nodes
|
||||
|
||||
mask_nodes = self.match_parent_path(add_or_sub_before_softmax,
|
||||
['Mul', 'Sub', 'Cast', 'Slice', 'Unsqueeze'],
|
||||
[ 1, 1, 1, 0, 0])
|
||||
if mask_nodes is not None:
|
||||
return mask_nodes
|
||||
|
||||
|
||||
mask_nodes = self.match_parent_path(add_or_sub_before_softmax,
|
||||
['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'],
|
||||
[ 1, None, 1, 0, 0])
|
||||
return mask_nodes
|
||||
|
||||
def check_attention_input(self, matmul_q, matmul_k, matmul_v, parent, output_name_to_node):
|
||||
reshape_nodes = []
|
||||
|
||||
for x in [matmul_q, matmul_k, matmul_v]:
|
||||
root_input = x.input[0]
|
||||
root_node = output_name_to_node[root_input]
|
||||
if root_node == parent:
|
||||
continue
|
||||
if root_node.op_type == 'Reshape' and root_node.input[0] == parent.output[0]:
|
||||
reshape_nodes.append(root_node)
|
||||
continue
|
||||
logger.debug(f"Check attention input failed:{root_input}, {parent.output[0]}")
|
||||
return False, []
|
||||
|
||||
return True, reshape_nodes
|
||||
|
||||
def fuse_attention(self):
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
|
@ -71,8 +106,10 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
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
|
||||
qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Div', '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,
|
||||
|
|
@ -94,28 +131,16 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
(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])
|
||||
mask_nodes = self.match_mask_path(qk_nodes[1])
|
||||
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("Failed to match mask path")
|
||||
continue
|
||||
if not self.has_constant_input(mask_nodes[1], 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:
|
||||
|
||||
is_same_root, reshape_nodes = self.check_attention_input(matmul_q, matmul_k, matmul_v, parent, output_name_to_node)
|
||||
if 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])
|
||||
|
|
@ -125,10 +150,8 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
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)
|
||||
nodes_to_remove.extend(reshape_nodes)
|
||||
attention_count += 1
|
||||
|
||||
nodes_to_remove.append(extra_reshape_0)
|
||||
self.replace_node_input(add, extra_reshape_0.output[0], matmul.output[0])
|
||||
else:
|
||||
|
|
@ -138,35 +161,46 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
self.update_graph()
|
||||
logger.info(f"Fused Attention count:{attention_count}")
|
||||
|
||||
def preprocess(self):
|
||||
self.process_embedding()
|
||||
self.fuse_mask()
|
||||
self.skip_reshape()
|
||||
|
||||
def skip_reshape(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
|
||||
|
||||
count = 0
|
||||
reshape_nodes = self.get_nodes_by_op_type("Reshape")
|
||||
for reshape_node in reshape_nodes:
|
||||
parent = self.get_parent(reshape_node, 0)
|
||||
if parent is None or parent.op_type == "Reshape":
|
||||
reshape_node.input[0] = parent.input[0]
|
||||
count += 1
|
||||
|
||||
if count > 0:
|
||||
logger.info(f"Skip consequent Reshape count: {count}")
|
||||
|
||||
def fuse_embedding(self, node, output_name_to_node):
|
||||
assert node.op_type == 'LayerNormalization'
|
||||
pos_embed_path2 = self.match_parent_path(
|
||||
logger.debug(f"start fusing embedding from node with output={node.output[0]}...")
|
||||
word_embed_path = 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))
|
||||
if word_embed_path is None:
|
||||
logger.debug("failed to match word_embed_path")
|
||||
return False
|
||||
|
||||
skip_node, add_node, gather_node = word_embed_path
|
||||
|
||||
word_initializer = self.get_initializer(gather_node.input[0])
|
||||
if word_initializer is None:
|
||||
logger.debug("failed to get word initializer")
|
||||
return False
|
||||
|
||||
temp = numpy_helper.to_array(word_initializer)
|
||||
|
|
@ -177,12 +211,50 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
logger.info("Failed to find word embedding. name:{}, shape:{}".format(word_initializer.name, temp.shape))
|
||||
return False
|
||||
|
||||
|
||||
pos_initializer = self.get_initializer(add_node.input[1])
|
||||
if pos_initializer is not None:
|
||||
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
|
||||
else:
|
||||
pos_embed_path = self.match_parent_path(
|
||||
add_node,
|
||||
['Gather', 'Slice'],
|
||||
[1, 1],
|
||||
output_name_to_node)
|
||||
if pos_embed_path is None:
|
||||
logger.debug("failed to match pos_embed_path")
|
||||
return False
|
||||
|
||||
pos_gather, pos_slice = pos_embed_path
|
||||
pos_initializer = self.get_initializer(pos_gather.input[0])
|
||||
if pos_initializer is None:
|
||||
logger.debug("failed to get pos initializer")
|
||||
return False
|
||||
|
||||
temp = numpy_helper.to_array(pos_initializer)
|
||||
if len(temp.shape) == 2:
|
||||
logger.info("Found word embedding. name:{}, shape:{}".format(pos_initializer.name, temp.shape))
|
||||
position_embedding = pos_initializer.name
|
||||
else:
|
||||
logger.info("Failed to find position embedding. name:{}, shape:{}".format(pos_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":
|
||||
logger.debug("failed to get gather")
|
||||
return False
|
||||
|
||||
segment_initializer = self.get_initializer(gather.input[0])
|
||||
if segment_initializer is None:
|
||||
logger.debug("failed to get segment initializer")
|
||||
return False
|
||||
|
||||
temp = numpy_helper.to_array(segment_initializer)
|
||||
|
|
@ -207,7 +279,60 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
if node.op_type == 'LayerNormalization':
|
||||
if self.fuse_embedding(node, output_name_to_node):
|
||||
return
|
||||
break
|
||||
|
||||
def fuse_mask(self):
|
||||
nodes_to_remove = []
|
||||
for node in self.nodes():
|
||||
if node.op_type == 'Mul' and self.has_constant_input(node, -10000):
|
||||
mask_path = self.match_parent_path(
|
||||
node,
|
||||
['Sub', 'Cast', 'Slice', 'Unsqueeze'],
|
||||
[ 0, 1, 0, 0])
|
||||
if mask_path is None:
|
||||
continue
|
||||
sub_node, cast_node, slice_node, unsqueeze_node = mask_path
|
||||
|
||||
mask_input_name = next(iter(self.mask_indice))
|
||||
if unsqueeze_node.input[0] != mask_input_name:
|
||||
print("Cast input {} is not mask input{}".format(unsqueeze_node.input[0], mask_input_name))
|
||||
continue
|
||||
|
||||
unsqueeze_added_1 = onnx.helper.make_node(
|
||||
'Unsqueeze',
|
||||
inputs=[mask_input_name],
|
||||
outputs=['mask_fuse_unsqueeze1_output'],
|
||||
name='Mask_UnSqueeze_1',
|
||||
axes=[1])
|
||||
|
||||
unsqueeze_added_2 = onnx.helper.make_node(
|
||||
'Unsqueeze',
|
||||
inputs=['mask_fuse_unsqueeze1_output'],
|
||||
outputs=['mask_fuse_unsqueeze2_output'],
|
||||
name='Mask_UnSqueeze_2',
|
||||
axes=[2])
|
||||
|
||||
#self.replace_node_input(cast_node, cast_node.input[0], 'mask_fuse_unsqueeze2_output')
|
||||
cast_node_2 = onnx.helper.make_node(
|
||||
'Cast',
|
||||
inputs=['mask_fuse_unsqueeze2_output'],
|
||||
outputs=['mask_fuse_cast_output'])
|
||||
cast_node_2.attribute.extend([onnx.helper.make_attribute("to", 1)])
|
||||
self.replace_node_input(sub_node, sub_node.input[1], 'mask_fuse_cast_output')
|
||||
|
||||
nodes_to_remove.extend([slice_node, unsqueeze_node, cast_node])
|
||||
self.add_node(unsqueeze_added_1)
|
||||
self.add_node(unsqueeze_added_2)
|
||||
self.add_node(cast_node_2)
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
|
||||
# Prune graph is done after removing nodes to remove island nodes.
|
||||
if len(nodes_to_remove) > 0:
|
||||
self.prune_graph()
|
||||
|
||||
logger.info("Fused mask" if len(nodes_to_remove) > 0 else "Failed to fuse mask")
|
||||
|
||||
def remove_extra_reshape(self):
|
||||
skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
|
||||
reshape_removed = 0
|
||||
|
|
@ -228,7 +353,7 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
self.remove_node(reshape_1)
|
||||
reshape_removed += 3
|
||||
|
||||
logger.info(f"Remove {reshape_removed} Reshape nodes.")
|
||||
return reshape_removed
|
||||
|
||||
def remove_extra_reshape_2(self):
|
||||
skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
|
||||
|
|
@ -257,9 +382,97 @@ class BertOnnxModelKeras(BertOnnxModelTF):
|
|||
|
||||
reshape_removed += 4
|
||||
|
||||
logger.info(f"Remove {reshape_removed} Reshape nodes.")
|
||||
return reshape_removed
|
||||
|
||||
def postprocess(self):
|
||||
self.remove_extra_reshape()
|
||||
self.remove_extra_reshape_2()
|
||||
reshape_removed = self.remove_extra_reshape() + self.remove_extra_reshape_2()
|
||||
logger.info(f"Remove {reshape_removed} Reshape nodes.")
|
||||
|
||||
self.prune_graph()
|
||||
|
||||
"""
|
||||
Fuse Gelu with Erf into one node:
|
||||
+------------------------------------------+
|
||||
| |
|
||||
| v
|
||||
[root] --> Div -----> Erf --> Add --> Mul -->Mul
|
||||
(B=1.4142...) (A=1) (A=0.5)
|
||||
|
||||
Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
|
||||
"""
|
||||
def fuse_gelu_with_elf(self, gelu_op_name):
|
||||
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 node in self.get_nodes_by_op_type('Erf'):
|
||||
erf_node = node
|
||||
|
||||
if erf_node.output[0] not in input_name_to_nodes:
|
||||
continue
|
||||
children = input_name_to_nodes[erf_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != 'Add':
|
||||
continue
|
||||
add_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(add_after_erf, 1):
|
||||
continue
|
||||
|
||||
if add_after_erf.output[0] not in input_name_to_nodes:
|
||||
continue
|
||||
children = input_name_to_nodes[add_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != 'Mul':
|
||||
continue
|
||||
mul_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(mul_after_erf, 0.5):
|
||||
continue
|
||||
|
||||
if mul_after_erf.output[0] not in input_name_to_nodes:
|
||||
continue
|
||||
children = input_name_to_nodes[mul_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != 'Mul':
|
||||
continue
|
||||
mul = children[0]
|
||||
|
||||
div = self.match_parent(erf_node, 'Div', 0, output_name_to_node)
|
||||
if div is None:
|
||||
continue
|
||||
|
||||
sqrt_node = None
|
||||
if self.find_constant_input(div, 1.4142, delta=0.001) != 1:
|
||||
sqrt_node = self.match_parent(div, 'Sqrt', 1, output_name_to_node)
|
||||
if sqrt_node is None:
|
||||
continue
|
||||
if not self.has_constant_input(sqrt_node, 2.0):
|
||||
continue
|
||||
|
||||
root_node = self.get_parent(div, 0, output_name_to_node)
|
||||
if root_node is None:
|
||||
continue
|
||||
|
||||
if root_node.output[0] not in mul.input:
|
||||
continue
|
||||
|
||||
subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul]
|
||||
if sqrt_node:
|
||||
subgraph_nodes.append(sqrt_node)
|
||||
|
||||
if not self.is_safe_to_fuse_nodes(subgraph_nodes, [mul.output[0]], input_name_to_nodes, output_name_to_node):
|
||||
continue
|
||||
|
||||
nodes_to_remove.extend(subgraph_nodes)
|
||||
gelu_node = onnx.helper.make_node(gelu_op_name,
|
||||
inputs=[root_node.output[0]],
|
||||
outputs=[mul.output[0]])
|
||||
gelu_node.domain = "com.microsoft"
|
||||
nodes_to_add.append(gelu_node)
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
self.add_nodes(nodes_to_add)
|
||||
if len(nodes_to_add) > 0:
|
||||
logger.info("Fused {} count:{}".format('FastGelu (approximation)' if gelu_op_name == 'FastGelu' else 'Gelu', len(nodes_to_add)))
|
||||
else:
|
||||
super().fuse_gelu_with_elf(gelu_op_name)
|
||||
|
|
|
|||
|
|
@ -119,3 +119,4 @@ python bert_perf_test.py --model optimized_model_gpu.onnx --batch_size 1 --seque
|
|||
```
|
||||
|
||||
After test is finished, a file like perf_results_CPU_B1_S128_<date_time>.txt or perf_results_GPU_B1_S128_<date_time>.txt will be output to the model directory.
|
||||
|
||||
|
|
|
|||
|
|
@ -92,10 +92,10 @@ def parse_arguments():
|
|||
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")
|
||||
help="number of attention heads. 12 for bert-base model and 16 for bert-large")
|
||||
|
||||
parser.add_argument('--hidden_size', required=False, type=int, default=768,
|
||||
help="bert model hidden size. 768 for base model and 1024 for large")
|
||||
help="bert model hidden size. 768 for bert-base model and 1024 for bert-large")
|
||||
|
||||
parser.add_argument('--sequence_length', required=False, type=int, default=128,
|
||||
help="max sequence length")
|
||||
|
|
|
|||
|
|
@ -208,4 +208,4 @@ def main():
|
|||
print("Test summary is saved to", summary_file)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
|
@ -23,7 +23,8 @@ 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_squad_pytorch1.4_opset10_fp32": 'test_data\\bert_squad_pytorch1.4_opset10_fp32\\BertForQuestionAnswering.onnx',
|
||||
"bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx'
|
||||
"bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx',
|
||||
"bert_keras_squad": 'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx'
|
||||
}
|
||||
|
||||
class TestBertOptimization(unittest.TestCase):
|
||||
|
|
@ -181,5 +182,14 @@ class TestBertOptimization(unittest.TestCase):
|
|||
}
|
||||
self.verify_node_count(bert_model, expected_node_count)
|
||||
|
||||
def test_keras_squad_model_cpu(self):
|
||||
input = BERT_TEST_MODELS['bert_keras_squad']
|
||||
|
||||
bert_model = optimize_model(input, 'bert_keras', gpu_only=False,
|
||||
num_heads=2, hidden_size=8, sequence_length=7,
|
||||
input_int32=False, float16=False)
|
||||
|
||||
self.assertTrue(bert_model.is_fully_optimized())
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -0,0 +1,291 @@
|
|||
#-------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
#--------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Convert a Bert large model exported by Keras2Onnx to a tiny model for test purpose.
|
||||
The input model is generated like the following (need install keras2onnx from source):
|
||||
|
||||
import numpy
|
||||
import keras2onnx
|
||||
from transformers import (TFBertForQuestionAnswering, BertTokenizer)
|
||||
|
||||
tokenizer = BertTokenizer.from_pretrained'bert-large-uncased-whole-word-masking-finetuned-squad', do_lower_case=True, cache_dir=cache_dir)
|
||||
model = TFBertForQuestionAnswering.from_pretrained(model_name_or_path, cache_dir=cache_dir)
|
||||
|
||||
question, text = "What is ONNX Runtime?", "ONNX Runtime is a performance-focused inference engine for ONNX models."
|
||||
inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors='tf')
|
||||
|
||||
output_model_path = os.path.join(output_dir, 'keras_{}.onnx'.format(model_name_or_path))
|
||||
if not os.path.exists(output_model_path):
|
||||
model.predict(inputs)
|
||||
onnx_model = keras2onnx.convert_keras(model, model.name)
|
||||
keras2onnx.save_model(onnx_model, output_model_path)
|
||||
"""
|
||||
|
||||
import onnx
|
||||
import onnx.utils
|
||||
import sys
|
||||
import argparse
|
||||
import numpy as np
|
||||
from onnx import ModelProto, TensorProto, numpy_helper
|
||||
from bert_model_optimization import OnnxModel
|
||||
import os
|
||||
import onnxruntime
|
||||
import random
|
||||
from pathlib import Path
|
||||
import timeit
|
||||
|
||||
DICT_SIZE = 20
|
||||
SEQ_LEN = 7
|
||||
|
||||
""" This class creates a tiny bert model for test purpose. """
|
||||
class TinyBertOnnxModel(OnnxModel):
|
||||
def __init__(self, model, verbose):
|
||||
super(TinyBertOnnxModel, self).__init__(model, verbose)
|
||||
self.resize_model()
|
||||
|
||||
def resize_weight(self, initializer_name, target_shape):
|
||||
weight = self.get_initializer(initializer_name)
|
||||
w = numpy_helper.to_array(weight)
|
||||
|
||||
target_w = w
|
||||
if len(target_shape) == 1:
|
||||
target_w = w[:target_shape[0]]
|
||||
elif len(target_shape) == 2:
|
||||
target_w = w[:target_shape[0],:target_shape[1]]
|
||||
elif len(target_shape) == 3:
|
||||
target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]]
|
||||
else:
|
||||
print("at most 3 dimensions")
|
||||
|
||||
tensor = onnx.helper.make_tensor(
|
||||
name=initializer_name + '_resize',
|
||||
data_type=TensorProto.FLOAT,
|
||||
dims=target_shape,
|
||||
vals=target_w.flatten().tolist())
|
||||
|
||||
return tensor
|
||||
|
||||
def resize_model(self):
|
||||
graph = self.model.graph
|
||||
initializers = graph.initializer
|
||||
|
||||
# parameters of input base model.
|
||||
old_parameters = {
|
||||
"seq_len" : 26,
|
||||
"hidden_size": 1024,
|
||||
"num_heads": 16,
|
||||
"size_per_head": 64,
|
||||
"word_dict_size": [28996, 30522], # list of supported dictionary size.
|
||||
"max_word_position": 512
|
||||
}
|
||||
|
||||
# parameters of output tiny model.
|
||||
new_parameters = {
|
||||
"seq_len" : SEQ_LEN,
|
||||
"hidden_size": 8,
|
||||
"num_heads": 2,
|
||||
"size_per_head": 4,
|
||||
"word_dict_size": DICT_SIZE,
|
||||
"max_word_position": 10
|
||||
}
|
||||
|
||||
for input in graph.input:
|
||||
if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]):
|
||||
print("input", input.name, input.type.tensor_type.shape)
|
||||
input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"]
|
||||
print("=>", input.type.tensor_type.shape)
|
||||
|
||||
reshapes = {}
|
||||
for initializer in initializers:
|
||||
tensor = numpy_helper.to_array(initializer)
|
||||
dtype = np.float32 if initializer.data_type == 1 else np.int32
|
||||
if len(tensor.shape) == 1 and tensor.shape[0] == 1:
|
||||
if tensor == old_parameters["num_heads"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]")
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["seq_len"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]")
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["size_per_head"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]")
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]")
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype), initializer.name))
|
||||
elif tensor == 4 * old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, 4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]")
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype), initializer.name))
|
||||
elif len(tensor.shape) == 0:
|
||||
if tensor == old_parameters["num_heads"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["num_heads"], "=>", new_parameters["num_heads"])
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["seq_len"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["seq_len"], "=>", new_parameters["seq_len"])
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["size_per_head"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["size_per_head"], "=>", new_parameters["size_per_head"])
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype), initializer.name))
|
||||
elif tensor == old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, old_parameters["hidden_size"], "=>", new_parameters["hidden_size"])
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype), initializer.name))
|
||||
elif tensor == 4 * old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, 4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"])
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype), initializer.name))
|
||||
elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]):
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name, 1.0 / np.sqrt(old_parameters["size_per_head"]), "=>", 1.0 / np.sqrt(new_parameters["size_per_head"]))
|
||||
initializer.CopyFrom(numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype), initializer.name))
|
||||
|
||||
new_shape=[]
|
||||
shape_changed = False
|
||||
for dim in tensor.shape:
|
||||
if (dim == old_parameters["hidden_size"]):
|
||||
new_shape.append(new_parameters["hidden_size"])
|
||||
shape_changed = True
|
||||
elif (dim == 4 * old_parameters["hidden_size"]):
|
||||
new_shape.append(4 * new_parameters["hidden_size"])
|
||||
shape_changed = True
|
||||
elif (dim in old_parameters["word_dict_size"]):
|
||||
new_shape.append(new_parameters["word_dict_size"])
|
||||
shape_changed = True
|
||||
elif (dim == old_parameters["max_word_position"]):
|
||||
new_shape.append(new_parameters["max_word_position"])
|
||||
shape_changed = True
|
||||
else:
|
||||
new_shape.append(dim)
|
||||
if shape_changed:
|
||||
reshapes[initializer.name] = new_shape
|
||||
print("initializer", initializer.name, tensor.shape, "=>", new_shape)
|
||||
|
||||
for initializer_name in reshapes:
|
||||
self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize')
|
||||
tensor = self.resize_weight(initializer_name, reshapes[initializer_name])
|
||||
self.model.graph.initializer.extend([tensor])
|
||||
|
||||
self.use_dynamic_axes()
|
||||
|
||||
def use_dynamic_axes(self, dynamic_batch_dim='batch_size', seq_len=7):
|
||||
"""
|
||||
Update input and output shape to use dynamic axes.
|
||||
"""
|
||||
dynamic_batch_inputs = {}
|
||||
for input in self.model.graph.input:
|
||||
dim_proto = input.type.tensor_type.shape.dim[0]
|
||||
dim_proto.dim_param = dynamic_batch_dim
|
||||
dim_proto = input.type.tensor_type.shape.dim[1]
|
||||
dim_proto.dim_value = 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
|
||||
dim_proto = output.type.tensor_type.shape.dim[1]
|
||||
dim_proto.dim_value = seq_len
|
||||
|
||||
def generate_test_data(onnx_file, output_path, batch_size, sequence_length, use_cpu = True, input_tensor_only=False, dictionary_size=DICT_SIZE, test_cases=3):
|
||||
input_data_type = np.int32
|
||||
for test_case in range(test_cases):
|
||||
input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=input_data_type)
|
||||
tensor_1 = numpy_helper.from_array(input_1, 'input_ids')
|
||||
|
||||
actual_seq_len = random.randint(sequence_length - 3, sequence_length)
|
||||
input_2 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
|
||||
temp = np.ones((batch_size, actual_seq_len), dtype=input_data_type)
|
||||
input_2[:temp.shape[0],:temp.shape[1]]=temp
|
||||
tensor_2 = numpy_helper.from_array(input_2, 'attention_mask')
|
||||
|
||||
input_3 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
|
||||
tensor_3 = numpy_helper.from_array(input_3, 'token_type_ids')
|
||||
|
||||
path = os.path.join(output_path, 'test_data_set_' + str(test_case))
|
||||
try:
|
||||
os.mkdir(path)
|
||||
except OSError:
|
||||
print ("Creation of the directory %s failed" % path)
|
||||
else:
|
||||
print ("Successfully created the directory %s " % path)
|
||||
|
||||
if input_tensor_only:
|
||||
return
|
||||
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider'])
|
||||
|
||||
input1_name = sess.get_inputs()[0].name
|
||||
output_names = [output.name for output in sess.get_outputs()]
|
||||
inputs = {'input_ids': input_1, 'attention_mask': input_2, 'token_type_ids': input_3}
|
||||
print ("inputs", inputs)
|
||||
result=sess.run(output_names, inputs)
|
||||
|
||||
with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
|
||||
f.write(tensor_1.SerializeToString())
|
||||
with open(os.path.join(path, 'input_{}.pb'.format(1)), 'wb') as f:
|
||||
f.write(tensor_2.SerializeToString())
|
||||
with open(os.path.join(path, 'input_{}.pb'.format(2)), 'wb') as f:
|
||||
f.write(tensor_3.SerializeToString())
|
||||
|
||||
for i, output_name in enumerate(output_names):
|
||||
tensor_result = numpy_helper.from_array(np.asarray(result[i]).reshape((batch_size, sequence_length)), output_names[i])
|
||||
with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f:
|
||||
f.write(tensor_result.SerializeToString())
|
||||
|
||||
start_time = timeit.default_timer()
|
||||
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
|
||||
|
||||
path_prefix = onnx_file[:-5] #remove .onnx suffix
|
||||
if use_cpu:
|
||||
sess_options.optimized_model_filepath = path_prefix + "_optimized_cpu.onnx"
|
||||
else:
|
||||
sess_options.optimized_model_filepath = path_prefix + "_optimized_gpu.onnx"
|
||||
|
||||
session = onnxruntime.InferenceSession(onnx_file, sess_options)
|
||||
if use_cpu:
|
||||
session.set_providers(['CPUExecutionProvider']) # use cpu
|
||||
else:
|
||||
if 'CUDAExecutionProvider' not in session.get_providers():
|
||||
print("Warning: GPU not found")
|
||||
continue
|
||||
outputs = session.run(None, inputs)
|
||||
evalTime = timeit.default_timer() - start_time
|
||||
if outputs[0].tolist() != result[0].tolist():
|
||||
print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format(use_cpu, result[0].tolist(), outputs[1].tolist()))
|
||||
print("** Evaluation done in total {} secs".format(evalTime))
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input', required=True, type=str)
|
||||
parser.add_argument('--output', required=True, type=str)
|
||||
parser.add_argument('--float16', required=False, action='store_true')
|
||||
parser.set_defaults(float16=False)
|
||||
args = parser.parse_args()
|
||||
|
||||
model = ModelProto()
|
||||
with open(args.input, "rb") as f:
|
||||
model.ParseFromString(f.read())
|
||||
|
||||
bert_model = TinyBertOnnxModel(model, False)
|
||||
|
||||
if args.float16:
|
||||
bert_model.convert_model_float32_to_float16()
|
||||
|
||||
bert_model.update_graph()
|
||||
bert_model.remove_unused_constant()
|
||||
|
||||
print("opset verion", bert_model.model.opset_import[0].version)
|
||||
|
||||
with open(args.output, "wb") as out:
|
||||
out.write(bert_model.model.SerializeToString())
|
||||
|
||||
p = Path(args.output)
|
||||
data_path = p.parent
|
||||
|
||||
batch_size = 1
|
||||
sequence_length = SEQ_LEN
|
||||
|
||||
generate_test_data(args.output, data_path, batch_size, sequence_length, use_cpu=not args.float16)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1 @@
|
|||
Boutput_1J 23= {<7B>=?є=N<><4E>=(,<2C>=\<5C><>=`<60><><
|
||||
|
|
@ -0,0 +1 @@
|
|||
Boutput_2J ÁŸ¼+2ª¼5à³¼`Çß¼…¹½2RÕ¼o‡;
|
||||
Loading…
Reference in a new issue