mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-05 04:17:53 +00:00
[transformers optimizer] catch symbolic shape inference exception and clean up (#7560)
catch symbolic shape inference exception. no prune graph when there is inner graph (Loop/If/Scan) add an wrapper for numpy_helper.to_array so that we can debug onnx graph without external data remove fuse_mask that is not used any more in onnx_model_bert_tf.py
This commit is contained in:
Tianlei Wu
GitHub
parent
faea7a222d
commit
3c9ece4a11
6 changed files with 51 additions and 125 deletions
|
|
@ -9,7 +9,7 @@ from typing import Tuple, Union
|
|||
from onnx import helper, numpy_helper, TensorProto, NodeProto
|
||||
from onnx_model import OnnxModel
|
||||
from fusion_base import Fusion
|
||||
from fusion_utils import FusionUtils
|
||||
from fusion_utils import FusionUtils, NumpyHelper
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
|
@ -107,7 +107,7 @@ class FusionAttention(Fusion):
|
|||
logger.debug(f"{reshape_q.input[1]} is not initializer.")
|
||||
return self.num_heads, self.hidden_size # Fall back to user specified value
|
||||
|
||||
q_shape_value = numpy_helper.to_array(q_shape)
|
||||
q_shape_value = NumpyHelper.to_array(q_shape)
|
||||
if len(q_shape_value) != 4 or (q_shape_value[2] <= 0 or q_shape_value[3] <= 0):
|
||||
logger.debug(f"q_shape_value={q_shape_value}. Expected value are like [0, 0, num_heads, head_size].")
|
||||
return self.num_heads, self.hidden_size # Fall back to user specified value
|
||||
|
|
@ -145,7 +145,7 @@ class FusionAttention(Fusion):
|
|||
Returns:
|
||||
Union[NodeProto, None]: the node created or None if failed.
|
||||
"""
|
||||
assert num_heads > 0 or hidden_size > 0 or (hidden_size % num_heads) == 0
|
||||
assert num_heads > 0 and hidden_size > 0 and (hidden_size % num_heads) == 0
|
||||
|
||||
q_weight = self.model.get_initializer(q_matmul.input[1])
|
||||
k_weight = self.model.get_initializer(k_matmul.input[1])
|
||||
|
|
@ -159,9 +159,9 @@ class FusionAttention(Fusion):
|
|||
return None
|
||||
if not (k_weight and v_weight and q_bias and k_bias):
|
||||
return None
|
||||
qw = numpy_helper.to_array(q_weight)
|
||||
kw = numpy_helper.to_array(k_weight)
|
||||
vw = numpy_helper.to_array(v_weight)
|
||||
qw = NumpyHelper.to_array(q_weight)
|
||||
kw = NumpyHelper.to_array(k_weight)
|
||||
vw = NumpyHelper.to_array(v_weight)
|
||||
|
||||
# Check if all matrices have the same shape
|
||||
assert qw.shape == kw.shape == vw.shape
|
||||
|
|
@ -173,9 +173,9 @@ class FusionAttention(Fusion):
|
|||
|
||||
qkv_weight = np.stack((qw, kw, vw), axis=1)
|
||||
|
||||
qb = numpy_helper.to_array(q_bias)
|
||||
kb = numpy_helper.to_array(k_bias)
|
||||
vb = numpy_helper.to_array(v_bias)
|
||||
qb = NumpyHelper.to_array(q_bias)
|
||||
kb = NumpyHelper.to_array(k_bias)
|
||||
vb = NumpyHelper.to_array(v_bias)
|
||||
|
||||
# 1d bias shape: [outsize,]. 2d bias shape: [a, b] where a*b = out_size
|
||||
assert qb.shape == kb.shape == vb.shape
|
||||
|
|
@ -196,7 +196,7 @@ class FusionAttention(Fusion):
|
|||
|
||||
# Sometimes weights and bias are stored in fp16
|
||||
if q_weight.data_type == 10:
|
||||
weight.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(weight).astype(np.float16), weight.name))
|
||||
weight.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(weight).astype(np.float16), weight.name))
|
||||
self.model.add_initializer(weight)
|
||||
|
||||
bias = helper.make_tensor(name=attention_node_name + '_qkv_bias',
|
||||
|
|
@ -204,7 +204,7 @@ class FusionAttention(Fusion):
|
|||
dims=[3 * out_size],
|
||||
vals=qkv_bias.flatten().tolist())
|
||||
if q_bias.data_type == 10:
|
||||
bias.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(bias).astype(np.float16), bias.name))
|
||||
bias.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(bias).astype(np.float16), bias.name))
|
||||
self.model.add_initializer(bias)
|
||||
|
||||
attention_inputs = [input, attention_node_name + '_qkv_weight', attention_node_name + '_qkv_bias']
|
||||
|
|
|
|||
|
|
@ -4,9 +4,10 @@
|
|||
#--------------------------------------------------------------------------
|
||||
|
||||
from logging import getLogger
|
||||
from onnx import helper, numpy_helper
|
||||
from onnx import helper
|
||||
from onnx_model import OnnxModel
|
||||
from fusion_base import Fusion
|
||||
from fusion_utils import NumpyHelper
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
|
@ -38,7 +39,7 @@ class FusionBiasGelu(Fusion):
|
|||
if initializer is None:
|
||||
continue
|
||||
bias_index = i
|
||||
bias_weight = numpy_helper.to_array(initializer)
|
||||
bias_weight = NumpyHelper.to_array(initializer)
|
||||
break
|
||||
if bias_weight is None:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -4,9 +4,10 @@
|
|||
#--------------------------------------------------------------------------
|
||||
|
||||
from logging import getLogger
|
||||
from onnx import helper, numpy_helper
|
||||
from onnx import helper
|
||||
from onnx_model import OnnxModel
|
||||
from fusion_base import Fusion
|
||||
from fusion_utils import NumpyHelper
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
|
@ -99,7 +100,7 @@ class FusionBiasSkipLayerNormalization(Fusion):
|
|||
if initializer is None:
|
||||
continue
|
||||
bias_index = i
|
||||
bias_weight = numpy_helper.to_array(initializer)
|
||||
bias_weight = NumpyHelper.to_array(initializer)
|
||||
break
|
||||
if bias_weight is None:
|
||||
logger.debug(f"Bias weight not found")
|
||||
|
|
|
|||
|
|
@ -3,9 +3,10 @@
|
|||
# Licensed under the MIT License.
|
||||
#--------------------------------------------------------------------------
|
||||
from logging import getLogger
|
||||
from onnx_model import OnnxModel
|
||||
from typing import Tuple
|
||||
from onnx import helper, TensorProto
|
||||
from onnx import helper, numpy_helper, TensorProto
|
||||
from numpy import ndarray
|
||||
from onnx_model import OnnxModel
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
|
@ -55,3 +56,14 @@ class FusionUtils:
|
|||
output_name = node.output[0]
|
||||
self.model.remove_node(node)
|
||||
self.model.replace_input_of_all_nodes(output_name, input_name)
|
||||
|
||||
class NumpyHelper:
|
||||
@staticmethod
|
||||
def to_array(tensor:TensorProto, fill_zeros:bool = False) -> ndarray:
|
||||
# When weights are in external data format but not presented, we can still test the optimizer with two changes:
|
||||
# (1) set fill_zeros = True (2) change load_external_data=False in optimizer.py
|
||||
if fill_zeros:
|
||||
from onnx import mapping
|
||||
return ndarray(shape=tensor.dims, dtype=mapping.TENSOR_TYPE_TO_NP_TYPE[tensor.data_type])
|
||||
|
||||
return numpy_helper.to_array(tensor)
|
||||
|
|
@ -32,9 +32,12 @@ class OnnxModel:
|
|||
if self.shape_infer_helper is None:
|
||||
self.shape_infer_helper = SymbolicShapeInferenceHelper(self.model)
|
||||
shape_infer_helper = self.shape_infer_helper
|
||||
|
||||
if shape_infer_helper.infer(dynamic_axis_mapping):
|
||||
return shape_infer_helper
|
||||
try:
|
||||
if shape_infer_helper.infer(dynamic_axis_mapping):
|
||||
return shape_infer_helper
|
||||
except:
|
||||
print("failed in shape inference", sys.exc_info()[0])
|
||||
|
||||
return None
|
||||
|
||||
def input_name_to_nodes(self):
|
||||
|
|
@ -585,6 +588,14 @@ class OnnxModel:
|
|||
Args:
|
||||
outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept.
|
||||
"""
|
||||
|
||||
for node in self.model.graph.node:
|
||||
# Some operators with inner graph in attributes like 'body' 'else_branch' or 'then_branch'
|
||||
if node.op_type in ['Loop', 'Scan', 'If']:
|
||||
# TODO: handle inner graph
|
||||
logger.debug(f"Skip prune_graph since graph has operator: {node.op_type}")
|
||||
return
|
||||
|
||||
if outputs is None:
|
||||
outputs = [output.name for output in self.model.graph.output]
|
||||
|
||||
|
|
@ -712,4 +723,4 @@ class OnnxModel:
|
|||
for input in self.model.graph.input:
|
||||
if self.get_initializer(input.name) is None:
|
||||
graph_inputs.append(input)
|
||||
return graph_inputs
|
||||
return graph_inputs
|
||||
|
|
@ -42,106 +42,9 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
|
||||
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 fuse_mask(self):
|
||||
nodes_to_remove = []
|
||||
for node in self.nodes():
|
||||
if node.op_type == 'Sub':
|
||||
parent_path_constant = self.match_parent_path(
|
||||
node,
|
||||
['Reshape', 'Mul', 'ConstantOfShape', 'Cast', 'Concat', 'Unsqueeze', 'Cast', 'Squeeze', 'Slice', 'Cast', 'Shape'],
|
||||
[ 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # yapf: disable
|
||||
if parent_path_constant is None:
|
||||
continue
|
||||
reshape_node_0, mul_node_0, constantofshape_node, cast_node_0, concat_node_0, unsqueeze_node, cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node = parent_path_constant
|
||||
|
||||
parent_path_mask = self.match_parent_path(
|
||||
mul_node_0,
|
||||
['Cast', 'Reshape', 'Cast', 'Concat', 'Unsqueeze'],
|
||||
[ 1, 0, 1, 0, 0]) # yapf: disable
|
||||
|
||||
if parent_path_mask is None:
|
||||
continue
|
||||
|
||||
cast_node_3, reshape_node_1, cast_node_4, concat_node_1, unsqueeze_node_1 = parent_path_mask
|
||||
|
||||
if not unsqueeze_node_1 == unsqueeze_node:
|
||||
continue
|
||||
|
||||
unsqueeze_added_1 = onnx.helper.make_node('Unsqueeze',
|
||||
inputs=[reshape_node_1.input[0]],
|
||||
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=[cast_node_3.input[0]],
|
||||
name='Mask_UnSqueeze_2',
|
||||
axes=[2])
|
||||
node.input[1] = cast_node_3.output[0]
|
||||
|
||||
nodes_to_remove.extend([
|
||||
reshape_node_0, mul_node_0, constantofshape_node, cast_node_0, concat_node_0, unsqueeze_node,
|
||||
cast_node_1, squeeze_node, slice_node, cast_node_2, shape_node
|
||||
])
|
||||
nodes_to_remove.extend([reshape_node_1, cast_node_4, concat_node_1])
|
||||
self.add_node(unsqueeze_added_1)
|
||||
self.add_node(unsqueeze_added_2)
|
||||
|
||||
self.remove_nodes(nodes_to_remove)
|
||||
if len(nodes_to_remove) > 0:
|
||||
logger.info("Fused mask")
|
||||
else:
|
||||
self.fuse_mask_2()
|
||||
|
||||
def fuse_mask_2(self):
|
||||
nodes_to_remove = []
|
||||
for node in self.nodes():
|
||||
if node.op_type == 'Mul' and self.has_constant_input(node, -10000):
|
||||
mask_path = self.match_parent_path(node, ['Sub', '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 = self.attention_mask.get_first_mask()
|
||||
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 get_2d_initializers_from_parent_subgraphs(self, current_node):
|
||||
"""
|
||||
Find initializers that is 2D. Returns a dictionary with name as key and shape as value.
|
||||
|
|
@ -432,6 +335,7 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
if q_nodes is None:
|
||||
logger.debug("Failed to match q path")
|
||||
continue
|
||||
|
||||
add_q = q_nodes[-2]
|
||||
matmul_q = q_nodes[-1]
|
||||
|
||||
|
|
@ -469,11 +373,12 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
if is_same_root:
|
||||
mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0])
|
||||
logger.debug("Create an Attention node.")
|
||||
|
||||
# For tf models, q and v are flipped.
|
||||
attention_node = self.attention_fusion.create_attention_node(mask_index, matmul_k, matmul_q, matmul_v,
|
||||
add_k, add_q, add_v, self.num_heads,
|
||||
self.hidden_size, parent.output[0],
|
||||
qkv_nodes[2].output[0])
|
||||
qkv_nodes[2].output[0])
|
||||
if attention_node is None:
|
||||
continue
|
||||
|
||||
|
|
@ -504,7 +409,6 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
self.add_initializer(tensor)
|
||||
parent.input[1] = parent.name + "_modified"
|
||||
|
||||
|
||||
self.add_node(attention_node)
|
||||
attention_count += 1
|
||||
|
||||
|
|
@ -524,8 +428,6 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
def preprocess(self):
|
||||
self.remove_identity()
|
||||
self.process_embedding()
|
||||
#TODO: remove fuse mask since we have embedding fused so fuse_attention shall handle the mask nodes.
|
||||
# self.fuse_mask()
|
||||
self.skip_reshape()
|
||||
|
||||
def skip_reshape(self):
|
||||
|
|
@ -554,5 +456,4 @@ class BertOnnxModelTF(BertOnnxModel):
|
|||
|
||||
def postprocess(self):
|
||||
self.remove_reshape_before_first_attention()
|
||||
# Temporary work around for the following comment as it will cause topological issues for a bert model
|
||||
# self.prune_graph()
|
||||
self.prune_graph()
|
||||
|
|
|
|||
Loading…
Reference in a new issue