mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Add GPT2 Attention Fusion in optimization script (#3488)
* Add Attention fusion for GPT2 * Support distilgpt2 in benchmark_gpt2.py * Add options to disable Attention/SkipLayerNormalization/EmbedLayerNormalization/BiasGelu fusions * Add logging at the begining of each fusion * Update notebooks: Add Gpt2OnnxModel.py to list of script files. * Add test for gpt2 model optimization * Add optional parameters (--input_ids --segment_ids --input_mask) for graph inputs * Fuse BiasGelu * Handle model that does not have segment_ids input. * Allow fuse embed layer without mask
This commit is contained in:
Tianlei Wu
GitHub
parent
5d3b217039
commit
7f46f347db
15 changed files with 967 additions and 151 deletions
|
|
@ -15,6 +15,19 @@ from OnnxModel import OnnxModel
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BertOptimizationOptions:
|
||||
|
||||
def __init__(self, model_type):
|
||||
self.enable_attention = True
|
||||
self.enable_skip_layer_norm = True
|
||||
self.enable_embed_layer_norm = True
|
||||
self.enable_bias_skip_layer_norm = True
|
||||
self.enable_bias_gelu = True
|
||||
|
||||
if model_type == 'gpt2':
|
||||
self.enable_skip_layer_norm = False
|
||||
|
||||
|
||||
class BertOnnxModel(OnnxModel):
|
||||
|
||||
def __init__(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
|
||||
|
|
@ -38,15 +51,31 @@ class BertOnnxModel(OnnxModel):
|
|||
|
||||
self.bert_inputs = []
|
||||
|
||||
def cast_input_to_int32(self, input_name):
|
||||
cast_output = input_name + '_int32'
|
||||
|
||||
# Avoid consequent Cast nodes.
|
||||
inputs = [input_name]
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
if input_name in output_name_to_node:
|
||||
parent_node = output_name_to_node[input_name]
|
||||
if parent_node and parent_node.op_type == 'Cast':
|
||||
inputs = [parent_node.input[0]]
|
||||
|
||||
cast_node = onnx.helper.make_node('Cast', inputs=inputs, outputs=[cast_output])
|
||||
cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.INT32))])
|
||||
self.add_node(cast_node)
|
||||
|
||||
return cast_output, cast_node
|
||||
|
||||
def cast_graph_input_to_int32(self, input_name):
|
||||
graph_input = self.find_graph_input(input_name)
|
||||
if graph_input is not None and graph_input.type.tensor_type.elem_type != TensorProto.INT32:
|
||||
cast_output = input_name + '_int32'
|
||||
cast_node = onnx.helper.make_node('Cast', inputs=[input_name], outputs=[cast_output])
|
||||
cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.INT32))])
|
||||
self.add_node(cast_node)
|
||||
cast_output, cast_node = self.cast_input_to_int32(input_name)
|
||||
logger.debug("Casted graph input {input_name} to int32")
|
||||
return True, cast_output
|
||||
|
||||
logger.debug(f"Did not cast graph input {input_name} to int32: found {graph_input is not None}")
|
||||
return False, input_name
|
||||
|
||||
def undo_cast_input_to_int32(self, input_name):
|
||||
|
|
@ -69,7 +98,12 @@ class BertOnnxModel(OnnxModel):
|
|||
return self.mask_indice[input]
|
||||
|
||||
# Add cast to convert int64 to int32
|
||||
casted, input_name = self.cast_graph_input_to_int32(input)
|
||||
if self.find_graph_input(input):
|
||||
casted, input_name = self.cast_graph_input_to_int32(input)
|
||||
else:
|
||||
input_name, cast_node = self.cast_input_to_int32(input)
|
||||
casted = True
|
||||
|
||||
if casted:
|
||||
self.mask_casted[input] = input_name
|
||||
|
||||
|
|
@ -125,19 +159,12 @@ class BertOnnxModel(OnnxModel):
|
|||
vals=qkv_weight.flatten().tolist())
|
||||
self.add_initializer(weight)
|
||||
|
||||
weight_input = onnx.helper.make_tensor_value_info(weight.name, TensorProto.FLOAT,
|
||||
[self.hidden_size, 3 * self.hidden_size])
|
||||
self.add_input(weight_input)
|
||||
|
||||
bias = onnx.helper.make_tensor(name=attention_node_name + '_qkv_bias',
|
||||
data_type=TensorProto.FLOAT,
|
||||
dims=[3 * self.hidden_size],
|
||||
vals=qkv_bias.flatten().tolist())
|
||||
self.add_initializer(bias)
|
||||
|
||||
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(
|
||||
'Attention',
|
||||
inputs=[input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias', mask_index],
|
||||
|
|
@ -350,6 +377,7 @@ class BertOnnxModel(OnnxModel):
|
|||
"""
|
||||
|
||||
def fuse_gelu_with_tanh(self):
|
||||
logger.debug(f"start FastGelu fusion...")
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
|
|
@ -435,19 +463,23 @@ class BertOnnxModel(OnnxModel):
|
|||
nodes_to_add.append(gelu_node)
|
||||
|
||||
if len(nodes_to_add) > 0:
|
||||
logger.info("Fused FastGelu count: {len(nodes_to_add)}")
|
||||
logger.info(f"Fused FastGelu count: {len(nodes_to_add)}")
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
self.add_nodes(nodes_to_add)
|
||||
|
||||
def fuse_add_bias_gelu(self):
|
||||
def fuse_bias_gelu(self, is_fastgelu):
|
||||
gelu_op_type = 'FastGelu' if is_fastgelu else 'Gelu'
|
||||
bias_gelu_op_type = 'FastGelu' if is_fastgelu else 'BiasGelu'
|
||||
logger.debug(f"start Bias and {gelu_op_type} fusion...")
|
||||
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
nodes_to_remove = []
|
||||
nodes_to_add = []
|
||||
|
||||
# Don't need to fuse Gelu+Add here because ORT native code can handle it
|
||||
for node in self.get_nodes_by_op_type('FastGelu'):
|
||||
for node in self.get_nodes_by_op_type(gelu_op_type):
|
||||
if len(node.input) != 1:
|
||||
continue
|
||||
|
||||
|
|
@ -476,20 +508,21 @@ class BertOnnxModel(OnnxModel):
|
|||
continue
|
||||
|
||||
nodes_to_remove.extend(subgraph_nodes)
|
||||
gelu_node = onnx.helper.make_node('FastGelu',
|
||||
gelu_node = onnx.helper.make_node(bias_gelu_op_type,
|
||||
inputs=[matmul.output[0], add.input[bias_index]],
|
||||
outputs=node.output,
|
||||
name=self.create_node_name('FastGelu', "FastGelu_AddBias_"))
|
||||
name=self.create_node_name(bias_gelu_op_type, gelu_op_type + "_AddBias_"))
|
||||
gelu_node.domain = "com.microsoft"
|
||||
nodes_to_add.append(gelu_node)
|
||||
|
||||
if len(nodes_to_add) > 0:
|
||||
logger.info(f"Fused FastGelu with Bias count:{len(nodes_to_add)}")
|
||||
logger.info(f"Fused {bias_gelu_op_type} with Bias count:{len(nodes_to_add)}")
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
self.add_nodes(nodes_to_add)
|
||||
|
||||
def fuse_add_bias_skip_layer_norm(self):
|
||||
logger.debug(f"start Bias and SkipLayerNormalization fusion...")
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
nodes_to_remove = []
|
||||
|
|
@ -552,6 +585,7 @@ class BertOnnxModel(OnnxModel):
|
|||
self.add_nodes(nodes_to_add)
|
||||
|
||||
def fuse_reshape(self):
|
||||
logger.debug(f"start Reshape fusion...")
|
||||
nodes = self.nodes()
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
|
@ -705,21 +739,12 @@ class BertOnnxModel(OnnxModel):
|
|||
SkipLayerNormalization
|
||||
"""
|
||||
|
||||
def fuse_embed_layer(self):
|
||||
def fuse_embed_layer_without_mask(self):
|
||||
logger.debug(f"start EmbedLayerNormalization (no mask) fusion...")
|
||||
nodes = self.nodes()
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
if len(self.mask_indice) == 0:
|
||||
logger.info("skip embed layer fusion since mask input is not found")
|
||||
return
|
||||
if len(self.mask_indice) > 1:
|
||||
logger.info("skip embed layer fusion since there are multiple mask inputs found")
|
||||
return
|
||||
mask_input_name = next(iter(self.mask_indice))
|
||||
mask_output_name = self.mask_indice[mask_input_name]
|
||||
mask_node = output_name_to_node[mask_output_name]
|
||||
|
||||
nodes_to_remove = []
|
||||
|
||||
# Find the first normalize node could be embedding layer.
|
||||
|
|
@ -731,6 +756,14 @@ class BertOnnxModel(OnnxModel):
|
|||
if self.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) is not None:
|
||||
normalize_node = node
|
||||
break
|
||||
# In case user disables attention fusion, check whether subgraph looks like Attention.
|
||||
if node.output[0] not in input_name_to_nodes:
|
||||
continue
|
||||
children = input_name_to_nodes[node.output[0]]
|
||||
children_types = sorted([child.op_type for child in children])
|
||||
if children_types == ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']:
|
||||
normalize_node = node
|
||||
break
|
||||
|
||||
if normalize_node is None:
|
||||
if len(self.get_nodes_by_op_type("EmbedLayerNormalization")) == 0:
|
||||
|
|
@ -780,7 +813,9 @@ class BertOnnxModel(OnnxModel):
|
|||
segment_ids = segment_embedding_gather.input[1]
|
||||
|
||||
if position_embedding_expand:
|
||||
subgraph_nodes = self.get_parent_subgraph_nodes(position_embedding_expand, [], output_name_to_node)
|
||||
input_parent = self.get_parent(position_embedding_shape, 0, output_name_to_node)
|
||||
subgraph_nodes = self.get_parent_subgraph_nodes(position_embedding_expand,
|
||||
[input_parent] if input_parent else [], output_name_to_node)
|
||||
nodes_to_remove.extend(subgraph_nodes)
|
||||
|
||||
nodes_to_remove.extend(word_embedding_path)
|
||||
|
|
@ -788,22 +823,37 @@ class BertOnnxModel(OnnxModel):
|
|||
nodes_to_remove.extend(segment_embedding_path)
|
||||
|
||||
nodes_to_remove.extend([normalize_node])
|
||||
nodes_to_remove.extend([mask_node])
|
||||
|
||||
# store inputs for further processing
|
||||
self.bert_inputs = [input_ids, segment_ids, mask_input_name]
|
||||
if self.find_graph_input(input_ids):
|
||||
self.bert_inputs = [input_ids, segment_ids] if self.find_graph_input(segment_ids) else [input_ids]
|
||||
|
||||
if not self.input_int32:
|
||||
# When mask has been casted to int32, use that casted one as input of embed layer norm.
|
||||
if mask_input_name in self.mask_casted:
|
||||
mask_input_name = self.mask_casted[mask_input_name]
|
||||
|
||||
# Cast input_ids and segment_ids to int32.
|
||||
casted, input_ids = self.cast_graph_input_to_int32(input_ids)
|
||||
|
||||
casted, segment_ids = self.cast_graph_input_to_int32(segment_ids)
|
||||
# Cast input_ids and segment_ids to int32.
|
||||
if self.find_graph_input(input_ids):
|
||||
if not self.input_int32:
|
||||
casted, input_ids = self.cast_graph_input_to_int32(input_ids)
|
||||
else:
|
||||
self.undo_cast_input_to_int32(mask_input_name)
|
||||
input_ids, input_ids_cast_node = self.cast_input_to_int32(input_ids)
|
||||
|
||||
if self.find_graph_input(segment_ids):
|
||||
if not self.input_int32:
|
||||
casted, segment_ids = self.cast_graph_input_to_int32(segment_ids)
|
||||
else:
|
||||
segment_ids, segment_ids_cast_node = self.cast_input_to_int32(segment_ids)
|
||||
|
||||
segment_id_path = self.match_parent_path(
|
||||
segment_ids_cast_node, ['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'],
|
||||
[0, 0, 1, 0, 0, 0])
|
||||
if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]:
|
||||
logger.debug("Simplify semgent id path...")
|
||||
self.add_node(
|
||||
onnx.helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"]))
|
||||
self.add_node(
|
||||
onnx.helper.make_node('ConstantOfShape',
|
||||
inputs=["input_shape"],
|
||||
outputs=["zeros_for_input_shape"],
|
||||
value=onnx.helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1])))
|
||||
segment_ids = "zeros_for_input_shape"
|
||||
|
||||
embed_node = onnx.helper.make_node(
|
||||
'EmbedLayerNormalization',
|
||||
|
|
@ -814,10 +864,9 @@ class BertOnnxModel(OnnxModel):
|
|||
position_embedding_weight_node.input[0],
|
||||
segment_embedding_gather.input[0],
|
||||
normalize_node.input[2],
|
||||
normalize_node.input[3], # gamma and beta
|
||||
mask_input_name
|
||||
normalize_node.input[3] # gamma and beta
|
||||
],
|
||||
outputs=["embed_output", mask_output_name],
|
||||
outputs=["embed_output", "dummy_mask_index"],
|
||||
name="EmbedLayer")
|
||||
|
||||
embed_node.domain = "com.microsoft"
|
||||
|
|
@ -826,13 +875,50 @@ class BertOnnxModel(OnnxModel):
|
|||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
self.add_node(embed_node)
|
||||
self.update_graph()
|
||||
logger.info("Fused EmbedLayerNormalization count: 1")
|
||||
self.prune_graph()
|
||||
|
||||
return embed_node
|
||||
|
||||
def fuse_embed_layer(self):
|
||||
embed_node = self.fuse_embed_layer_without_mask()
|
||||
if embed_node is None:
|
||||
logger.info("Fused EmbedLayerNormalization count: 0")
|
||||
return
|
||||
|
||||
if len(self.mask_indice) > 1:
|
||||
logger.info("There are multiple mask inputs found!")
|
||||
|
||||
if len(self.mask_indice) != 1:
|
||||
logger.info("Fused EmbedLayerNormalization (no mask) count: 1")
|
||||
return
|
||||
|
||||
mask_input_name = next(iter(self.mask_indice))
|
||||
mask_output_name = self.mask_indice[mask_input_name]
|
||||
mask_node = output_name_to_node[mask_output_name]
|
||||
|
||||
nodes_to_remove = []
|
||||
nodes_to_remove.extend([mask_node])
|
||||
|
||||
# store inputs for further processing
|
||||
self.bert_inputs.append(mask_input_name)
|
||||
|
||||
if not self.input_int32:
|
||||
# When mask has been casted to int32, use that casted one as input of embed layer norm.
|
||||
if mask_input_name in self.mask_casted:
|
||||
mask_input_name = self.mask_casted[mask_input_name]
|
||||
else:
|
||||
self.undo_cast_input_to_int32(mask_input_name)
|
||||
|
||||
embed_node.input[7] = mask_input_name
|
||||
embed_node.output[1] = mask_output_name
|
||||
logger.info("Added mask to EmbedLayerNormalization")
|
||||
|
||||
# Change graph input data type int32 if needed.
|
||||
if self.input_int32:
|
||||
self.change_input_to_int32()
|
||||
|
||||
logger.info("Fused EmbedLayerNormalization count: 1")
|
||||
|
||||
def get_bert_inputs(self, include_mask=True):
|
||||
return self.bert_inputs if include_mask else self.bert_inputs[:2]
|
||||
|
||||
|
|
@ -856,14 +942,13 @@ class BertOnnxModel(OnnxModel):
|
|||
graph = self.graph()
|
||||
|
||||
batch_size = self.get_batch_size_from_graph_input()
|
||||
input_batch_size = batch_size if isinstance(batch_size, int) else 1
|
||||
new_graph_inputs = []
|
||||
|
||||
bert_inputs = self.get_bert_inputs()
|
||||
for input in graph.input:
|
||||
if input.name in bert_inputs:
|
||||
int32_input = onnx.helper.make_tensor_value_info(input.name, TensorProto.INT32,
|
||||
[input_batch_size, self.sequence_length])
|
||||
input_shape = [batch_size if isinstance(batch_size, int) else 1, self.sequence_length]
|
||||
int32_input = onnx.helper.make_tensor_value_info(input.name, TensorProto.INT32, input_shape)
|
||||
new_graph_inputs.append(int32_input)
|
||||
else:
|
||||
new_graph_inputs.append(input)
|
||||
|
|
@ -924,6 +1009,8 @@ class BertOnnxModel(OnnxModel):
|
|||
| |
|
||||
+----------------------+
|
||||
"""
|
||||
logger.debug(f"start LayerNormalization fusion...")
|
||||
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
|
|
@ -1013,6 +1100,8 @@ class BertOnnxModel(OnnxModel):
|
|||
"""
|
||||
Fuse Add + LayerNormalization into one node: SkipLayerNormalization
|
||||
"""
|
||||
logger.debug(f"start SkipLayerNormaliation fusion...")
|
||||
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
|
|
@ -1043,9 +1132,9 @@ class BertOnnxModel(OnnxModel):
|
|||
return
|
||||
|
||||
def postprocess(self):
|
||||
return
|
||||
self.prune_graph()
|
||||
|
||||
def optimize(self):
|
||||
def optimize(self, options: BertOptimizationOptions = None):
|
||||
self.fuse_layer_norm()
|
||||
|
||||
self.fuse_gelu()
|
||||
|
|
@ -1054,20 +1143,28 @@ class BertOnnxModel(OnnxModel):
|
|||
|
||||
self.fuse_reshape()
|
||||
|
||||
self.fuse_skip_layer_norm()
|
||||
if (options is None) or options.enable_skip_layer_norm:
|
||||
self.fuse_skip_layer_norm()
|
||||
|
||||
self.fuse_attention()
|
||||
if (options is None) or options.enable_attention:
|
||||
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()
|
||||
if (options is None) or options.enable_embed_layer_norm:
|
||||
self.fuse_embed_layer()
|
||||
|
||||
# Post-processing like removing extra reshape nodes.
|
||||
self.postprocess()
|
||||
|
||||
# Bias fusion is done after postprocess to avoid extra Reshape between bias and Gelu/FastGelu/SkipLayerNormalization
|
||||
if (options is None) or options.enable_bias_gelu:
|
||||
# Fuse Gelu and Add Bias before it.
|
||||
self.fuse_bias_gelu(is_fastgelu=True)
|
||||
self.fuse_bias_gelu(is_fastgelu=False)
|
||||
|
||||
if (options is None) or options.enable_bias_skip_layer_norm:
|
||||
# Fuse SkipLayerNormalization and Add Bias before it.
|
||||
self.fuse_add_bias_skip_layer_norm()
|
||||
|
||||
if self.float16:
|
||||
self.convert_model_float32_to_float16()
|
||||
|
||||
|
|
@ -1090,6 +1187,7 @@ class BertOnnxModel(OnnxModel):
|
|||
for op in ops:
|
||||
nodes = self.get_nodes_by_op_type(op)
|
||||
op_count[op] = len(nodes)
|
||||
logger.info(f"Optimized operators:{op_count}")
|
||||
return op_count
|
||||
|
||||
def is_fully_optimized(self):
|
||||
|
|
|
|||
|
|
@ -2,11 +2,189 @@
|
|||
# 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 BertOnnxModel import BertOnnxModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Gpt2OnnxModel(BertOnnxModel):
|
||||
|
||||
def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
|
||||
super().__init__(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
|
||||
|
||||
def fuse_attention(self):
|
||||
"""
|
||||
Fuse Attention subgraph into one Attention node.
|
||||
"""
|
||||
logger.debug(f"start attention fusion...")
|
||||
|
||||
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_nodes_by_op_type("LayerNormalization"):
|
||||
return_indice = []
|
||||
qkv_nodes = self.match_parent_path(
|
||||
normalize_node,
|
||||
['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'],
|
||||
[0, None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice
|
||||
) # yapf: disable
|
||||
if qkv_nodes is None:
|
||||
continue
|
||||
(add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes
|
||||
|
||||
another_input = add_qkv.input[1 - return_indice[0]]
|
||||
|
||||
v_nodes = self.match_parent_path(
|
||||
matmul_qkv,
|
||||
['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'],
|
||||
[1, 0, 0, 0, 0, 0]) # yapf: disable
|
||||
if v_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match v path")
|
||||
continue
|
||||
(transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes
|
||||
|
||||
layernorm_before_attention = self.get_parent(reshape_before_gemm, 0, output_name_to_node)
|
||||
if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization':
|
||||
logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}")
|
||||
continue
|
||||
|
||||
if not another_input in layernorm_before_attention.input:
|
||||
logger.debug("Add and LayerNormalization shall have one same input")
|
||||
continue
|
||||
|
||||
qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0])
|
||||
if qk_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match qk path")
|
||||
continue
|
||||
(softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
|
||||
|
||||
q_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0])
|
||||
if q_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match q path")
|
||||
continue
|
||||
(transpose_q, reshape_q, split_q) = q_nodes
|
||||
if split_v != split_q:
|
||||
logger.debug("fuse_attention: skip since split_v != split_q")
|
||||
continue
|
||||
|
||||
k_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0])
|
||||
if k_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match k path")
|
||||
continue
|
||||
(transpose_k, reshape_k, split_k) = k_nodes
|
||||
if split_v != split_k:
|
||||
logger.debug("fuse_attention: skip since split_v != split_k")
|
||||
continue
|
||||
|
||||
mask_nodes = self.match_parent_path(
|
||||
sub_qk,
|
||||
['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
|
||||
[1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
|
||||
if mask_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match mask path")
|
||||
continue
|
||||
(mul_mask, sub_mask, slice_mask, slice_mask_0, unsqueeze_mask, sub_mask, squeeze_mask, slice_mask_1,
|
||||
shape_mask, div_mask) = mask_nodes
|
||||
|
||||
if div_qk != div_mask:
|
||||
logger.debug("fuse_attention: skip since div_qk != div_mask")
|
||||
continue
|
||||
|
||||
self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0],
|
||||
attention_count == 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)
|
||||
attention_count += 1
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
self.prune_graph()
|
||||
logger.info(f"Fused Attention count:{attention_count}")
|
||||
|
||||
def create_attention_node(self, gemm, gemm_qkv, input, output, add_graph_input):
|
||||
attention_node_name = self.create_node_name('Attention')
|
||||
attention_node = onnx.helper.make_node('Attention',
|
||||
inputs=[input, gemm.input[1], gemm.input[2]],
|
||||
outputs=[attention_node_name + "_output"],
|
||||
name=attention_node_name)
|
||||
attention_node.domain = "com.microsoft"
|
||||
attention_node.attribute.extend(
|
||||
[onnx.helper.make_attribute("num_heads", self.num_heads),
|
||||
onnx.helper.make_attribute("unidirectional", 1)])
|
||||
|
||||
matmul_node = onnx.helper.make_node('MatMul',
|
||||
inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
|
||||
outputs=[attention_node_name + "_matmul_output"],
|
||||
name=attention_node_name + "_matmul")
|
||||
|
||||
add_node = onnx.helper.make_node('Add',
|
||||
inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]],
|
||||
outputs=[output],
|
||||
name=attention_node_name + "_add")
|
||||
|
||||
self.add_node(attention_node)
|
||||
self.add_node(matmul_node)
|
||||
self.add_node(add_node)
|
||||
|
||||
def postprocess(self):
|
||||
"""
|
||||
Remove extra reshape nodes.
|
||||
"""
|
||||
logger.debug(f"start postprocessing...")
|
||||
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
reshape_count = 0
|
||||
for gemm_node in self.get_nodes_by_op_type("Gemm"):
|
||||
reshape_after_gemm = self.find_first_child_by_type(gemm_node,
|
||||
'Reshape',
|
||||
input_name_to_nodes,
|
||||
recursive=False)
|
||||
|
||||
return_indice = []
|
||||
nodes = self.match_parent_path(gemm_node, ['Reshape', 'FastGelu'], [0, 0], output_name_to_node)
|
||||
if nodes is None:
|
||||
nodes = self.match_parent_path(gemm_node, ['Reshape', 'LayerNormalization'], [0, 0],
|
||||
output_name_to_node)
|
||||
if nodes is None:
|
||||
continue
|
||||
(reshape_before_gemm, root_node) = nodes
|
||||
|
||||
matmul_node_name = self.create_node_name('MatMul', 'FullyConnect_MatMul')
|
||||
matmul_node = onnx.helper.make_node('MatMul',
|
||||
inputs=[matmul_node_name + "_input", gemm_node.input[1]],
|
||||
outputs=[matmul_node_name + "_output"],
|
||||
name=matmul_node_name)
|
||||
|
||||
add_node_name = self.create_node_name('Add', 'FullyConnect_Add')
|
||||
add_node = onnx.helper.make_node('Add',
|
||||
inputs=[matmul_node_name + "_output", gemm_node.input[2]],
|
||||
outputs=[add_node_name + "_output"],
|
||||
name=add_node_name)
|
||||
|
||||
root_node.output[0] = matmul_node_name + "_input"
|
||||
self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output")
|
||||
|
||||
self.add_node(matmul_node)
|
||||
self.add_node(add_node)
|
||||
|
||||
reshape_count += 2
|
||||
|
||||
self.prune_graph()
|
||||
logger.info(f"Remove Reshape count:{reshape_count}")
|
||||
|
|
|
|||
|
|
@ -11,10 +11,16 @@ import psutil
|
|||
import argparse
|
||||
import logging
|
||||
import torch
|
||||
from transformers import GPT2Model, GPT2Tokenizer
|
||||
from transformers import GPT2Model, GPT2LMHeadModel, GPT2Tokenizer
|
||||
|
||||
logger = logging.getLogger('')
|
||||
|
||||
# Map alias to a tuple of Model Class and pretrained model name
|
||||
MODEL_CLASSES = {
|
||||
"gpt2": (GPT2Model, GPT2Tokenizer, "gpt2"),
|
||||
"distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2")
|
||||
}
|
||||
|
||||
|
||||
def dump_environment():
|
||||
if "OMP_NUM_THREADS" in os.environ:
|
||||
|
|
@ -44,9 +50,7 @@ def pytorch_inference(model, input_ids, past=None, total_runs=100):
|
|||
with torch.no_grad():
|
||||
for _ in range(total_runs):
|
||||
start = time.time()
|
||||
outputs = model(
|
||||
input_ids=input_ids,
|
||||
past=past) #attention_mask=inputs['attention_mask'], token_type_ids=inputs['token_type_ids'],
|
||||
outputs = model(input_ids=input_ids, past=past)
|
||||
latency.append(time.time() - start)
|
||||
|
||||
logger.info("PyTorch Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f')))
|
||||
|
|
@ -88,6 +92,12 @@ def inference(model, ort_session, input_ids, past=None, total_runs=100, verify_o
|
|||
def parse_arguments():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--model_type',
|
||||
required=True,
|
||||
type=str,
|
||||
choices=list(MODEL_CLASSES.keys()),
|
||||
help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
|
||||
|
||||
parser.add_argument('--cache_dir', required=True, type=str, help="cache directory")
|
||||
|
||||
parser.add_argument('--output_dir', required=True, type=str, help="output onnx model directory")
|
||||
|
|
@ -132,6 +142,24 @@ def setup_logger(verbose=True):
|
|||
logger.setLevel(logging_level)
|
||||
|
||||
|
||||
def remove_past_outputs(export_model_path):
|
||||
from onnx import ModelProto
|
||||
from OnnxModel import OnnxModel
|
||||
|
||||
model = ModelProto()
|
||||
with open(export_model_path, "rb") as f:
|
||||
model.ParseFromString(f.read())
|
||||
bert_model = OnnxModel(model)
|
||||
|
||||
# remove past state outputs and only keep the first output.
|
||||
keep_output_names = [bert_model.model.graph.output[0].name]
|
||||
logger.info(f"Prune graph to keep the first output and drop past state outputs:{keep_output_names}")
|
||||
bert_model.prune_graph(keep_output_names)
|
||||
onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_out1.onnx'.format(int(enable_past_input)))
|
||||
bert_model.save_model_to_file(onnx_model_path)
|
||||
return onnx_model_path
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_arguments()
|
||||
setup_logger(args.verbose)
|
||||
|
|
@ -147,7 +175,8 @@ def main():
|
|||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
model_class, tokenizer_class, model_name_or_path = (GPT2Model, GPT2Tokenizer, 'gpt2')
|
||||
(model_class, tokenizer_class, model_name_or_path) = MODEL_CLASSES[args.model_type]
|
||||
|
||||
tokenizer = tokenizer_class.from_pretrained(model_name_or_path, cache_dir=cache_dir)
|
||||
model = model_class.from_pretrained(model_name_or_path, cache_dir=cache_dir)
|
||||
model.eval().cpu()
|
||||
|
|
@ -197,10 +226,11 @@ def main():
|
|||
setup_environment(args.use_openmp)
|
||||
import onnxruntime
|
||||
|
||||
onnx_model_path = export_model_path
|
||||
onnx_model_path = export_model_path if enable_past_input else remove_past_outputs(export_model_path)
|
||||
|
||||
if args.enable_optimization:
|
||||
from bert_model_optimization import optimize_model
|
||||
m = optimize_model(export_model_path,
|
||||
m = optimize_model(onnx_model_path,
|
||||
model_type='gpt2',
|
||||
gpu_only=False,
|
||||
num_heads=12,
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ import argparse
|
|||
import numpy as np
|
||||
from collections import deque
|
||||
from onnx import ModelProto, TensorProto, numpy_helper
|
||||
from BertOnnxModel import BertOnnxModel
|
||||
from BertOnnxModel import BertOnnxModel, BertOptimizationOptions
|
||||
from BertOnnxModelTF import BertOnnxModelTF
|
||||
from BertOnnxModelKeras import BertOnnxModelKeras
|
||||
from Gpt2OnnxModel import Gpt2OnnxModel
|
||||
|
|
@ -137,6 +137,33 @@ def parse_arguments():
|
|||
help="whether the target device is gpu or not")
|
||||
parser.set_defaults(gpu_only=False)
|
||||
|
||||
parser.add_argument('--disable_attention', required=False, action='store_true', help="disable Attention fusion")
|
||||
parser.set_defaults(disable_attention=False)
|
||||
|
||||
parser.add_argument('--disable_skip_layer_norm',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help="disable SkipLayerNormalization fusion")
|
||||
parser.set_defaults(disable_skip_layer_norm=False)
|
||||
|
||||
parser.add_argument('--disable_embed_layer_norm',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help="disable EmbedLayerNormalization fusion")
|
||||
parser.set_defaults(disable_embed_layer_norm=False)
|
||||
|
||||
parser.add_argument('--disable_bias_skip_layer_norm',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help="disable Add Bias and SkipLayerNormalization fusion")
|
||||
parser.set_defaults(disable_bias_skip_layer_norm=False)
|
||||
|
||||
parser.add_argument('--disable_bias_gelu',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help="disable Add Bias and Gelu/FastGelu fusion")
|
||||
parser.set_defaults(disable_bias_gelu=False)
|
||||
|
||||
parser.add_argument('--verbose', required=False, action='store_true')
|
||||
parser.set_defaults(verbose=False)
|
||||
|
||||
|
|
@ -152,6 +179,21 @@ def parse_arguments():
|
|||
return args
|
||||
|
||||
|
||||
def get_optimization_options(args):
|
||||
optimization_options = BertOptimizationOptions(args.model_type)
|
||||
if args.disable_attention:
|
||||
optimization_options.enable_attention = False
|
||||
if args.disable_skip_layer_norm:
|
||||
optimization_options.enable_skip_layer_norm = False
|
||||
if args.disable_embed_layer_norm:
|
||||
optimization_options.enable_embed_layer_norm = False
|
||||
if args.disable_bias_skip_layer_norm:
|
||||
optimization_options.enable_bias_skip_layer_norm = False
|
||||
if args.disable_bias_gelu:
|
||||
optimization_options.enable_bias_gelu = False
|
||||
return optimization_options
|
||||
|
||||
|
||||
def optimize_model(input,
|
||||
model_type,
|
||||
gpu_only,
|
||||
|
|
@ -160,7 +202,8 @@ def optimize_model(input,
|
|||
sequence_length,
|
||||
input_int32,
|
||||
float16,
|
||||
opt_level=99):
|
||||
opt_level=99,
|
||||
optimization_options=None):
|
||||
(optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type]
|
||||
|
||||
input_model_path = input
|
||||
|
|
@ -177,8 +220,11 @@ def optimize_model(input,
|
|||
f"Model producer not matched: Expect {producer}, Got {model.producer_name} {model.producer_version}. Please specify correct --model_type parameter."
|
||||
)
|
||||
|
||||
if optimization_options is None:
|
||||
optimization_options = BertOptimizationOptions(model_type)
|
||||
|
||||
bert_model = optimizer_class(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
|
||||
bert_model.optimize()
|
||||
bert_model.optimize(optimization_options)
|
||||
|
||||
return bert_model
|
||||
|
||||
|
|
@ -195,8 +241,11 @@ def main():
|
|||
|
||||
setup_logger(args.verbose)
|
||||
|
||||
optimization_options = get_optimization_options(args)
|
||||
|
||||
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, args.opt_level)
|
||||
args.sequence_length, args.input_int32, args.float16, args.opt_level,
|
||||
optimization_options)
|
||||
|
||||
bert_model.save_model_to_file(args.output)
|
||||
|
||||
|
|
|
|||
|
|
@ -240,9 +240,10 @@ def run_perf_tests(perf_results, model_path, batch_size, sequence_length, use_gp
|
|||
|
||||
|
||||
def run_performance(perf_results, model_path, batch_size, sequence_length, use_gpu, test_cases, test_times, seed,
|
||||
verbose, inclusive, test_all, no_warmup, opt_level):
|
||||
# Try deduce input names from model.
|
||||
input_ids, segment_ids, input_mask = get_bert_inputs(model_path)
|
||||
verbose, inclusive, test_all, no_warmup, opt_level, input_ids_name, segment_ids_name,
|
||||
input_mask_name):
|
||||
|
||||
input_ids, segment_ids, input_mask = get_bert_inputs(model_path, input_ids_name, segment_ids_name, input_mask_name)
|
||||
|
||||
# Do not generate random mask for performance test.
|
||||
print(f"Generating {test_cases} samples for batch_size={batch_size} sequence_length={sequence_length}")
|
||||
|
|
@ -347,6 +348,10 @@ def parse_arguments():
|
|||
parser.add_argument('--no_warmup', required=False, action='store_true', help="do not use one sample for warm-up.")
|
||||
parser.set_defaults(no_warmup=False)
|
||||
|
||||
parser.add_argument('--input_ids', required=False, type=str, default=None, help="input name for input ids")
|
||||
parser.add_argument('--segment_ids', required=False, type=str, default=None, help="input name for segment ids")
|
||||
parser.add_argument('--input_mask', required=False, type=str, default=None, help="input name for attention mask")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
|
@ -367,7 +372,7 @@ def main():
|
|||
for batch_size in batch_size_set:
|
||||
run_performance(perf_results, args.model, batch_size, args.sequence_length, args.use_gpu, args.samples,
|
||||
args.test_times, args.seed, args.verbose, args.inclusive, args.all, args.no_warmup,
|
||||
args.opt_level)
|
||||
args.opt_level, args.input_ids, args.segment_ids, args.input_mask)
|
||||
|
||||
# Sort the results so that the first one has smallest latency.
|
||||
sorted_results = sorted(perf_results.items(), reverse=False, key=lambda x: x[1])
|
||||
|
|
|
|||
|
|
@ -107,15 +107,22 @@ def fake_test_data(batch_size, sequence_length, test_cases, dictionary_size, ver
|
|||
"""
|
||||
Generate fake input data for test.
|
||||
"""
|
||||
assert input_ids is not None
|
||||
|
||||
np.random.seed(random_seed)
|
||||
random.seed(random_seed)
|
||||
|
||||
all_inputs = []
|
||||
for test_case in range(test_cases):
|
||||
input_1 = fake_input_ids_data(input_ids, batch_size, sequence_length, dictionary_size)
|
||||
input_2 = fake_segment_ids_data(segment_ids, batch_size, sequence_length)
|
||||
input_3 = fake_input_mask_data(input_mask, batch_size, sequence_length, random_mask_length)
|
||||
inputs = {input_ids.name: input_1, segment_ids.name: input_2, input_mask.name: input_3}
|
||||
inputs = {input_ids.name: input_1}
|
||||
|
||||
if segment_ids:
|
||||
inputs[segment_ids.name] = fake_segment_ids_data(segment_ids, batch_size, sequence_length)
|
||||
|
||||
if input_mask:
|
||||
inputs[input_mask.name] = fake_input_mask_data(input_mask, batch_size, sequence_length, random_mask_length)
|
||||
|
||||
if verbose and len(all_inputs) == 0:
|
||||
print("Example inputs", inputs)
|
||||
all_inputs.append(inputs)
|
||||
|
|
@ -144,7 +151,7 @@ def get_graph_input_from_embed_node(onnx_model, embed_node, input_index):
|
|||
return graph_input
|
||||
|
||||
|
||||
def get_bert_inputs(onnx_file):
|
||||
def get_bert_inputs(onnx_file, input_ids_name=None, segment_ids_name=None, input_mask_name=None):
|
||||
"""
|
||||
Get graph inputs for bert model.
|
||||
First, we will deduce from EmbedLayerNormalization node. If not found, we will guess based on naming.
|
||||
|
|
@ -154,8 +161,31 @@ def get_bert_inputs(onnx_file):
|
|||
model.ParseFromString(f.read())
|
||||
|
||||
onnx_model = OnnxModel(model)
|
||||
|
||||
graph_inputs = onnx_model.get_graph_inputs_excluding_initializers()
|
||||
|
||||
if input_ids_name is not None:
|
||||
input_ids = onnx_model.find_graph_input(input_ids_name)
|
||||
if input_ids is None:
|
||||
raise ValueError(f"Graph does not have input named {input_ids_name}")
|
||||
|
||||
segment_ids = None
|
||||
if segment_ids_name:
|
||||
segment_ids = onnx_model.find_graph_input(segment_ids_name)
|
||||
if segment_ids is None:
|
||||
raise ValueError(f"Graph does not have input named {segment_ids_name}")
|
||||
|
||||
input_mask = None
|
||||
if input_mask_name:
|
||||
input_mask = onnx_model.find_graph_input(input_mask_name)
|
||||
if input_mask is None:
|
||||
raise ValueError(f"Graph does not have input named {input_mask_name}")
|
||||
|
||||
expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0)
|
||||
if len(graph_inputs) != expected_inputs:
|
||||
raise ValueError(f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}")
|
||||
|
||||
return input_ids, segment_ids, input_mask
|
||||
|
||||
if len(graph_inputs) != 3:
|
||||
raise ValueError("Expect the graph to have 3 inputs. Got {}".format(len(graph_inputs)))
|
||||
|
||||
|
|
|
|||
|
|
@ -78,9 +78,11 @@ def compare(baseline_results, treatment_results, verbose, rtol=1e-3, atol=1e-4):
|
|||
|
||||
|
||||
def run_test(baseline_model, optimized_model, output_dir, batch_size, sequence_length, use_gpu, test_cases, seed,
|
||||
use_openmp, verbose, rtol, atol):
|
||||
use_openmp, verbose, rtol, atol, input_ids_name, segment_ids_name, input_mask_name):
|
||||
|
||||
# Try deduce input names from optimized model.
|
||||
input_ids, segment_ids, input_mask = get_bert_inputs(optimized_model)
|
||||
input_ids, segment_ids, input_mask = get_bert_inputs(optimized_model, input_ids_name, segment_ids_name,
|
||||
input_mask_name)
|
||||
|
||||
# Use random mask length for accuracy test. It might introduce slight inflation in latency reported in this script.
|
||||
all_inputs = generate_test_data(batch_size,
|
||||
|
|
@ -161,6 +163,10 @@ def parse_arguments():
|
|||
parser.add_argument('--verbose', required=False, action='store_true', help="print verbose information")
|
||||
parser.set_defaults(verbose=False)
|
||||
|
||||
parser.add_argument('--input_ids', required=False, type=str, default=None, help="input name for input ids")
|
||||
parser.add_argument('--segment_ids', required=False, type=str, default=None, help="input name for segment ids")
|
||||
parser.add_argument('--input_mask', required=False, type=str, default=None, help="input name for attention mask")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
|
@ -174,7 +180,8 @@ def main():
|
|||
path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
run_test(args.baseline_model, args.optimized_model, args.output_dir, args.batch_size, args.sequence_length,
|
||||
args.use_gpu, args.samples, args.seed, args.openmp, args.verbose, args.rtol, args.atol)
|
||||
args.use_gpu, args.samples, args.seed, args.openmp, args.verbose, args.rtol, args.atol, args.input_ids,
|
||||
args.segment_ids, args.input_mask)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -484,25 +484,25 @@
|
|||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
|
||||
"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"if not os.path.exists(script_dir):\n",
|
||||
" os.makedirs(script_dir)\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
|
||||
"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"if not os.path.exists(script_dir):\n",
|
||||
" os.makedirs(script_dir)\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@
|
|||
"source": [
|
||||
"In this tutorial, you'll be introduced to how to load a Bert model from PyTorch, convert it to ONNX, and inference it for high performance using ONNX Runtime and NVIDIA GPU. In the following sections, we are going to use the Bert model trained with Stanford Question Answering Dataset (SQuAD) dataset as an example. Bert SQuAD model is used in question answering scenarios, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable.\n",
|
||||
"\n",
|
||||
"This notebook is for CPU inference. For GPU inference, please look at another notebook [Inference PyTorch Bert Model with ONNX Runtime on CPU](PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb)."
|
||||
"This notebook is for GPU inference. For CPU inference, please look at another notebook [Inference PyTorch Bert Model with ONNX Runtime on CPU](PyTorch_Bert-Squad_OnnxRuntime_CPU.ipynb)."
|
||||
]
|
||||
},
|
||||
{
|
||||
|
|
@ -603,25 +603,25 @@
|
|||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
|
||||
"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"if not os.path.exists(script_dir):\n",
|
||||
" os.makedirs(script_dir)\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
|
||||
"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"if not os.path.exists(script_dir):\n",
|
||||
" os.makedirs(script_dir)\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
|
|
|||
|
|
@ -100,30 +100,30 @@
|
|||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"cache_dir = \"./squad\"\n",
|
||||
"output_dir = \"./output\"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"\n",
|
||||
"for directory in [cache_dir, output_dir, script_dir]:\n",
|
||||
" if not os.path.exists(directory):\n",
|
||||
" os.makedirs(directory)\n",
|
||||
"\n",
|
||||
"# Download scripts for BERT optimization.\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'OnnxModel.py', 'bert_model_optimization.py']\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import wget\n",
|
||||
"\n",
|
||||
"cache_dir = \"./squad\"\n",
|
||||
"output_dir = \"./output\"\n",
|
||||
"script_dir = './bert_scripts'\n",
|
||||
"\n",
|
||||
"for directory in [cache_dir, output_dir, script_dir]:\n",
|
||||
" if not os.path.exists(directory):\n",
|
||||
" os.makedirs(directory)\n",
|
||||
"\n",
|
||||
"# Download scripts for BERT optimization.\n",
|
||||
"url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
|
||||
"script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py']\n",
|
||||
"\n",
|
||||
"for filename in script_files:\n",
|
||||
" target_file = os.path.join(script_dir, filename)\n",
|
||||
" if enable_overwrite and os.path.exists(target_file):\n",
|
||||
" os.remove(target_file)\n",
|
||||
" if not os.path.exists(target_file):\n",
|
||||
" wget.download(url_prfix + filename, target_file)\n",
|
||||
" print(\"Downloaded\", filename)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
|
|
|||
|
|
@ -29,7 +29,9 @@ BERT_TEST_MODELS = {
|
|||
"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'
|
||||
'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx',
|
||||
"gpt2":
|
||||
'test_data\\gpt2_pytorch1.4_opset11_no_past\\GPT2Model.onnx'
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -202,8 +204,8 @@ class TestBertOptimization(unittest.TestCase):
|
|||
'Attention': 12,
|
||||
'LayerNormalization': 0,
|
||||
'SkipLayerNormalization': 24,
|
||||
'BiasGelu': 0,
|
||||
'Gelu': 12,
|
||||
'BiasGelu': 12,
|
||||
'Gelu': 0,
|
||||
'FastGelu': 0
|
||||
}
|
||||
self.verify_node_count(bert_model, expected_node_count)
|
||||
|
|
@ -222,6 +224,28 @@ class TestBertOptimization(unittest.TestCase):
|
|||
|
||||
self.assertTrue(bert_model.is_fully_optimized())
|
||||
|
||||
def test_gpt2(self):
|
||||
input = BERT_TEST_MODELS['gpt2']
|
||||
bert_model = optimize_model(input,
|
||||
'gpt2',
|
||||
gpu_only=False,
|
||||
num_heads=2,
|
||||
hidden_size=4,
|
||||
sequence_length=2,
|
||||
input_int32=False,
|
||||
float16=False)
|
||||
|
||||
expected_node_count = {
|
||||
'EmbedLayerNormalization': 0,
|
||||
'Attention': 12,
|
||||
'Gelu': 0,
|
||||
'FastGelu': 12,
|
||||
'BiasGelu': 0,
|
||||
'LayerNormalization': 25,
|
||||
'SkipLayerNormalization': 0
|
||||
}
|
||||
self.verify_node_count(bert_model, expected_node_count)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -0,0 +1,394 @@
|
|||
#-------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
#--------------------------------------------------------------------------
|
||||
# This tool generates a tiny GPT2 model for testing fusion script.
|
||||
# You can use benchmark_gpt2.py to get a gpt2 ONNX model as input of this tool.
|
||||
|
||||
import onnx
|
||||
import onnx.utils
|
||||
import sys
|
||||
import argparse
|
||||
import numpy as np
|
||||
from onnx import ModelProto, TensorProto, numpy_helper
|
||||
from OnnxModel import OnnxModel
|
||||
import os
|
||||
import onnxruntime
|
||||
import random
|
||||
from pathlib import Path
|
||||
import timeit
|
||||
|
||||
DICT_SIZE = 20
|
||||
SEQ_LEN = 2
|
||||
""" This class creates a tiny bert model for test purpose. """
|
||||
|
||||
# parameters of input base model.
|
||||
old_parameters = {
|
||||
"seq_len": 5,
|
||||
"hidden_size": 768,
|
||||
"num_heads": 12,
|
||||
"size_per_head": 64,
|
||||
"word_dict_size": [50257], # list of supported dictionary size.
|
||||
"max_word_position": 1024
|
||||
}
|
||||
|
||||
# parameters of output tiny model.
|
||||
new_parameters = {
|
||||
"seq_len": SEQ_LEN,
|
||||
"hidden_size": 4,
|
||||
"num_heads": 2,
|
||||
"size_per_head": 2,
|
||||
"word_dict_size": DICT_SIZE,
|
||||
"max_word_position": 8
|
||||
}
|
||||
|
||||
|
||||
class TinyBertOnnxModel(OnnxModel):
|
||||
|
||||
def __init__(self, model):
|
||||
super(TinyBertOnnxModel, self).__init__(model)
|
||||
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]]
|
||||
elif len(target_shape) == 4:
|
||||
target_w = w[:target_shape[0], :target_shape[1], :target_shape[2], :target_shape[3]]
|
||||
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
|
||||
|
||||
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)
|
||||
if initializer.data_type == TensorProto.FLOAT:
|
||||
dtype = np.float32
|
||||
elif initializer.data_type == TensorProto.INT32:
|
||||
dtype = np.int32
|
||||
elif initializer.data_type == TensorProto.INT64:
|
||||
dtype = np.int64
|
||||
else:
|
||||
print("data type not supported by this tool:", dtype)
|
||||
|
||||
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 tensor == 3 * old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name,
|
||||
3 * old_parameters["hidden_size"], "=>[", 3 * new_parameters["hidden_size"], "]")
|
||||
initializer.CopyFrom(
|
||||
numpy_helper.from_array(np.asarray([3 * 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 == 3 * old_parameters["hidden_size"]:
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name,
|
||||
3 * old_parameters["hidden_size"], "=>", 3 * new_parameters["hidden_size"])
|
||||
initializer.CopyFrom(
|
||||
numpy_helper.from_array(np.asarray(3 * 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))
|
||||
elif tensor == np.sqrt(old_parameters["size_per_head"]):
|
||||
print("initializer type={}".format(initializer.data_type), initializer.name,
|
||||
np.sqrt(old_parameters["size_per_head"]), "=>", np.sqrt(new_parameters["size_per_head"]))
|
||||
initializer.CopyFrom(
|
||||
numpy_helper.from_array(np.asarray(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 == 3 * old_parameters["hidden_size"]):
|
||||
new_shape.append(3 * 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])
|
||||
|
||||
# Add node name, replace split node attribute.
|
||||
nodes_to_add = []
|
||||
nodes_to_remove = []
|
||||
for i, node in enumerate(graph.node):
|
||||
if node.op_type == "Split":
|
||||
nodes_to_add.append(
|
||||
onnx.helper.make_node('Split',
|
||||
node.input,
|
||||
node.output,
|
||||
name="Split_{}".format(i),
|
||||
axis=2,
|
||||
split=[
|
||||
new_parameters["hidden_size"], new_parameters["hidden_size"],
|
||||
new_parameters["hidden_size"]
|
||||
]))
|
||||
nodes_to_remove.append(node)
|
||||
print("update split",
|
||||
[new_parameters["hidden_size"], new_parameters["hidden_size"], new_parameters["hidden_size"]])
|
||||
if node.op_type == "Constant":
|
||||
for att in node.attribute:
|
||||
if att.name == 'value':
|
||||
if numpy_helper.to_array(att.t) == old_parameters["num_heads"]:
|
||||
nodes_to_add.append(
|
||||
onnx.helper.make_node('Constant',
|
||||
inputs=node.input,
|
||||
outputs=node.output,
|
||||
value=onnx.helper.make_tensor(name=att.t.name,
|
||||
data_type=TensorProto.INT64,
|
||||
dims=[],
|
||||
vals=[new_parameters["num_heads"]
|
||||
])))
|
||||
print("constant", att.t.name, old_parameters["num_heads"], "=>",
|
||||
new_parameters["num_heads"])
|
||||
if numpy_helper.to_array(att.t) == np.sqrt(old_parameters["size_per_head"]):
|
||||
nodes_to_add.append(
|
||||
onnx.helper.make_node('Constant',
|
||||
inputs=node.input,
|
||||
outputs=node.output,
|
||||
value=onnx.helper.make_tensor(
|
||||
name=att.t.name,
|
||||
data_type=TensorProto.FLOAT,
|
||||
dims=[],
|
||||
vals=[np.sqrt(new_parameters["size_per_head"])])))
|
||||
print("constant", att.t.name, np.sqrt(old_parameters["size_per_head"]), "=>",
|
||||
np.sqrt(new_parameters["size_per_head"]))
|
||||
else:
|
||||
node.name = node.op_type + "_" + str(i)
|
||||
for node in nodes_to_remove:
|
||||
graph.node.remove(node)
|
||||
graph.node.extend(nodes_to_add)
|
||||
|
||||
def remove_past_outputs(self):
|
||||
keep_output_names = [self.model.graph.output[0].name] # remove past state outputs which is not needed.
|
||||
print(f"Prune graph to keep the first output and drop past state outputs:{keep_output_names}")
|
||||
self.prune_graph(keep_output_names)
|
||||
|
||||
|
||||
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=1,
|
||||
output_optimized_model=False):
|
||||
|
||||
input_data_type = np.int64
|
||||
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')
|
||||
|
||||
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 = {input1_name: input_1}
|
||||
result = sess.run(output_names, inputs)
|
||||
|
||||
with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
|
||||
f.write(tensor_1.SerializeToString())
|
||||
|
||||
for i, output_name in enumerate(output_names):
|
||||
if i == 0:
|
||||
tensor_result = numpy_helper.from_array(
|
||||
np.asarray(result[i]).reshape((batch_size, sequence_length, new_parameters["hidden_size"])),
|
||||
output_names[i])
|
||||
with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f:
|
||||
f.write(tensor_result.SerializeToString())
|
||||
else:
|
||||
tensor_result = numpy_helper.from_array(
|
||||
np.asarray(result[i]).reshape(
|
||||
(2, batch_size, new_parameters["num_heads"], sequence_length, new_parameters["size_per_head"])),
|
||||
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
|
||||
|
||||
if output_optimized_model:
|
||||
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 not np.allclose(outputs[0], result[0], rtol=1e-04, atol=1e-05):
|
||||
print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format(
|
||||
use_cpu, result[0].tolist(), outputs[0].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)
|
||||
parser.add_argument('--no_past_outputs', required=False, action='store_true')
|
||||
parser.set_defaults(no_past_outputs=False)
|
||||
parser.add_argument('--output_optimized_model', required=False, action='store_true')
|
||||
parser.set_defaults(output_optimized_model=False)
|
||||
args = parser.parse_args()
|
||||
|
||||
model = ModelProto()
|
||||
with open(args.input, "rb") as f:
|
||||
model.ParseFromString(f.read())
|
||||
|
||||
bert_model = TinyBertOnnxModel(model)
|
||||
|
||||
if args.float16:
|
||||
bert_model.convert_model_float32_to_float16()
|
||||
|
||||
if args.no_past_outputs:
|
||||
bert_model.remove_past_outputs()
|
||||
|
||||
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,
|
||||
output_optimized_model=args.output_optimized_model)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
|
@ -0,0 +1 @@
|
|||
B
hidden_statesJ ÙaÊ>ÛÑ>&ÏIÀ¢?½;g>v,²>©3CÀYŽ²?
|
||||
Loading…
Reference in a new issue