mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Improve symbolic shape inference in transformers tools (#12217)
improve symbolic shape inference handling n transformers tools: avoid infinite loop and suppress duplicated warnings
This commit is contained in:
Tianlei Wu
GitHub
parent
975bb56e8c
commit
972e5e7300
4 changed files with 96 additions and 71 deletions
|
|
@ -139,7 +139,6 @@ class SymbolicShapeInference:
|
|||
"Gather": self._infer_Gather,
|
||||
"GatherElements": self._infer_GatherElements,
|
||||
"GatherND": self._infer_GatherND,
|
||||
"Gelu": self._pass_on_shape_and_type,
|
||||
"Identity": self._pass_on_shape_and_type,
|
||||
"If": self._infer_If,
|
||||
"Loop": self._infer_Loop,
|
||||
|
|
@ -1536,7 +1535,7 @@ class SymbolicShapeInference:
|
|||
handle_negative_axis(ax, self._get_shape_rank(node, i + num_scan_states))
|
||||
for i, ax in enumerate(scan_input_axes)
|
||||
]
|
||||
# We may have cases where the subgraph has optionial inputs that appear in both subgraph's input and initializer,
|
||||
# We may have cases where the subgraph has optional inputs that appear in both subgraph's input and initializer,
|
||||
# but not in the node's input. In such cases, the input model might be invalid, but let's skip those optional inputs.
|
||||
assert len(subgraph.input) >= len(node.input)
|
||||
subgraph_inputs = subgraph.input[: len(node.input)]
|
||||
|
|
|
|||
|
|
@ -24,6 +24,10 @@ class FusionSkipLayerNormalization(Fusion):
|
|||
# Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
|
||||
self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True)
|
||||
|
||||
if self.shape_infer_helper is None:
|
||||
# TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op.
|
||||
logger.warning("symbolic shape inference disabled or failed.")
|
||||
|
||||
def fuse(self, node, input_name_to_nodes, output_name_to_node):
|
||||
add = self.model.get_parent(node, 0, output_name_to_node)
|
||||
|
||||
|
|
@ -43,15 +47,14 @@ class FusionSkipLayerNormalization(Fusion):
|
|||
if self.shape_infer_helper is not None:
|
||||
if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]):
|
||||
logger.debug(
|
||||
f"skip skiplayernorm fusion since shape of inputs ({add.input[0]}, {add.input[1]}) are not same"
|
||||
"skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same",
|
||||
add.input[0],
|
||||
add.input[1],
|
||||
)
|
||||
return
|
||||
else:
|
||||
# shape_infer_helper can not handle subgraphs. Current work around is to disable skiplayernorm fusion
|
||||
# longterm todo: support subgraph in symbolic_shape_infer or support add broadcasting in skiplayernorm op
|
||||
logger.warning(
|
||||
"symbolic shape infer failed. it's safe to ignore this message if there is no issue with optimized model"
|
||||
)
|
||||
logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed")
|
||||
return
|
||||
|
||||
gather_path = self.model.match_parent_path(add, ["Gather"], [None])
|
||||
if gather_path is not None and self.model.find_graph_input(gather_path[0].input[1]) is None:
|
||||
|
|
|
|||
|
|
@ -3,10 +3,10 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import onnx
|
||||
from typing import Dict
|
||||
|
||||
# In ORT Package the symbolic_shape_infer.py is in ../tools
|
||||
file_path = os.path.dirname(__file__)
|
||||
|
|
@ -17,85 +17,104 @@ else:
|
|||
|
||||
from symbolic_shape_infer import SymbolicShapeInference, get_shape_from_type_proto, sympy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SymbolicShapeInferenceHelper(SymbolicShapeInference):
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
verbose=0,
|
||||
int_max=2**31 - 1,
|
||||
auto_merge=True,
|
||||
guess_output_rank=False,
|
||||
):
|
||||
def __init__(self, model, verbose=0, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False):
|
||||
super().__init__(int_max, auto_merge, guess_output_rank, verbose)
|
||||
self.model_ = onnx.ModelProto()
|
||||
self.model_.CopyFrom(model)
|
||||
self.all_shapes_inferred_ = False
|
||||
self.inferred_ = False
|
||||
self.model_ = model
|
||||
self.all_shapes_inferred_: bool = False
|
||||
self.is_inferred_: bool = False
|
||||
self.dynamic_axis_mapping_: Dict[str, int] = {}
|
||||
|
||||
# The goal is to remove dynamic_axis_mapping
|
||||
def infer(self, dynamic_axis_mapping):
|
||||
if self.inferred_:
|
||||
def infer(self, dynamic_axis_mapping: Dict[str, int], max_runs: int = 128):
|
||||
"""Run shape inference, and try replace dynamic axis from string to integer when mapping is provided.
|
||||
|
||||
Args:
|
||||
dynamic_axis_mapping (_type_): a dictionary with name of dynamic axis as key, like {"batch_size" : 4}
|
||||
max_runs (int, optional): limit maximum number of runs to avoid infinite loop. Defaults to 32.
|
||||
|
||||
Returns:
|
||||
bool: whether all shapes has been inferred or not.
|
||||
"""
|
||||
assert dynamic_axis_mapping is not None
|
||||
|
||||
if self.is_inferred_ and self.dynamic_axis_mapping_ == dynamic_axis_mapping:
|
||||
return self.all_shapes_inferred_
|
||||
|
||||
self.dynamic_axis_mapping_ = dynamic_axis_mapping # e.g {"batch_size" : 4, "seq_len" :7}
|
||||
self.dynamic_axis_mapping_ = dynamic_axis_mapping
|
||||
|
||||
self._preprocess(self.model_)
|
||||
while self.run_:
|
||||
self.all_shapes_inferred_ = self._infer_impl()
|
||||
|
||||
self.inferred_ = True
|
||||
count = 0
|
||||
while self.run_:
|
||||
logger.debug(f"shape infer run {count}")
|
||||
self.all_shapes_inferred_ = self._infer_impl()
|
||||
count += 1
|
||||
if max_runs > 0 and count >= max_runs:
|
||||
break
|
||||
|
||||
self.is_inferred_ = True
|
||||
return self.all_shapes_inferred_
|
||||
|
||||
# override _preprocess() to avoid unnecessary model copy since ctor copies the model
|
||||
def _preprocess(self, in_mp):
|
||||
self.out_mp_ = in_mp
|
||||
self.graph_inputs_ = dict([(i.name, i) for i in list(self.out_mp_.graph.input)])
|
||||
self.initializers_ = dict([(i.name, i) for i in self.out_mp_.graph.initializer])
|
||||
self.known_vi_ = dict([(i.name, i) for i in list(self.out_mp_.graph.input)])
|
||||
self.known_vi_.update(
|
||||
dict(
|
||||
[
|
||||
(
|
||||
i.name,
|
||||
onnx.helper.make_tensor_value_info(i.name, i.data_type, list(i.dims)),
|
||||
)
|
||||
for i in self.out_mp_.graph.initializer
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Override _get_sympy_shape() in symbolic_shape_infer.py to ensure shape inference by giving the actual value of dynamic axis
|
||||
def _get_sympy_shape(self, node, idx):
|
||||
"""Override it to ensure shape inference by giving the actual value of dynamic axis."""
|
||||
sympy_shape = []
|
||||
for d in self._get_shape(node, idx):
|
||||
if type(d) == str:
|
||||
if d in self.dynamic_axis_mapping_.keys():
|
||||
sympy_shape.append(self.dynamic_axis_mapping_[d])
|
||||
elif d in self.symbolic_dims_:
|
||||
sympy_shape.append(self.symbolic_dims_[d])
|
||||
|
||||
shape = self._get_shape(node, idx)
|
||||
if shape:
|
||||
for dim in shape:
|
||||
if isinstance(dim, str):
|
||||
if dim in self.dynamic_axis_mapping_:
|
||||
sympy_shape.append(self.dynamic_axis_mapping_[dim])
|
||||
elif dim in self.symbolic_dims_:
|
||||
sympy_shape.append(self.symbolic_dims_[dim])
|
||||
else:
|
||||
sympy_shape.append(sympy.Symbol(dim, integer=True))
|
||||
else:
|
||||
sympy_shape.append(sympy.Symbol(d, integer=True))
|
||||
else:
|
||||
assert None != d
|
||||
sympy_shape.append(d)
|
||||
assert dim is not None
|
||||
sympy_shape.append(dim)
|
||||
return sympy_shape
|
||||
|
||||
def get_edge_shape(self, edge):
|
||||
assert self.all_shapes_inferred_ == True
|
||||
"""Get shape of an edge.
|
||||
|
||||
Args:
|
||||
edge (str): name of edge
|
||||
|
||||
Returns:
|
||||
Optional[List[int]]: the shape, or None if shape is unknown
|
||||
"""
|
||||
assert self.all_shapes_inferred_
|
||||
if edge not in self.known_vi_:
|
||||
print("Cannot retrive the shape of " + str(edge))
|
||||
print("Cannot retrieve the shape of " + str(edge))
|
||||
return None
|
||||
|
||||
type_proto = self.known_vi_[edge].type
|
||||
shape = get_shape_from_type_proto(type_proto)
|
||||
for i in range(len(shape)):
|
||||
d = shape[i]
|
||||
if type(d) == str and d in self.dynamic_axis_mapping_.keys():
|
||||
shape[i] = self.dynamic_axis_mapping_[d]
|
||||
|
||||
if shape is not None:
|
||||
for i, dim in enumerate(shape):
|
||||
if isinstance(dim, str) and dim in self.dynamic_axis_mapping_:
|
||||
shape[i] = self.dynamic_axis_mapping_[dim]
|
||||
|
||||
return shape
|
||||
|
||||
def compare_shape(self, edge, edge_other):
|
||||
assert self.all_shapes_inferred_ == True
|
||||
"""Compare shape of two edges.
|
||||
|
||||
Args:
|
||||
edge (str): name of edge
|
||||
edge_other (str): name of another edge
|
||||
|
||||
Raises:
|
||||
Exception: At least one shape is missed for edges to compare
|
||||
|
||||
Returns:
|
||||
bool: whether the shape is same or not
|
||||
"""
|
||||
assert self.all_shapes_inferred_
|
||||
shape = self.get_edge_shape(edge)
|
||||
shape_other = self.get_edge_shape(edge_other)
|
||||
if shape is None or shape_other is None:
|
||||
|
|
|
|||
|
|
@ -1,15 +1,17 @@
|
|||
import unittest
|
||||
|
||||
import onnx
|
||||
import pytest
|
||||
import torch
|
||||
from parity_utilities import find_transformers_source
|
||||
|
||||
if find_transformers_source():
|
||||
from benchmark_helper import OptimizerInfo, Precision
|
||||
from benchmark_helper import ConfigModifier, OptimizerInfo, Precision
|
||||
from huggingface_models import MODELS
|
||||
from onnx_exporter import export_onnx_model_from_pt
|
||||
from shape_infer_helper import SymbolicShapeInferenceHelper
|
||||
else:
|
||||
from onnxruntime.transformers.benchmark_helper import OptimizerInfo, Precision
|
||||
from onnxruntime.transformers.benchmark_helper import ConfigModifier, OptimizerInfo, Precision
|
||||
from onnxruntime.transformers.huggingface_models import MODELS
|
||||
from onnxruntime.transformers.onnx_exporter import export_onnx_model_from_pt
|
||||
from onnxruntime.transformers.shape_infer_helper import SymbolicShapeInferenceHelper
|
||||
|
|
@ -19,15 +21,18 @@ class SymbolicShapeInferenceHelperTest(unittest.TestCase):
|
|||
def _load_onnx(self, model_name):
|
||||
input_names = MODELS[model_name][0]
|
||||
base_path = "../onnx_models/"
|
||||
import torch
|
||||
|
||||
config_modifier = ConfigModifier(None)
|
||||
fusion_options = None
|
||||
model_class = "AutoModel"
|
||||
with torch.no_grad():
|
||||
export_onnx_model_from_pt(
|
||||
model_name,
|
||||
MODELS[model_name][1],
|
||||
MODELS[model_name][2],
|
||||
MODELS[model_name][3],
|
||||
None,
|
||||
model_class,
|
||||
config_modifier,
|
||||
"../cache_models",
|
||||
base_path,
|
||||
input_names[:1],
|
||||
|
|
@ -38,10 +43,9 @@ class SymbolicShapeInferenceHelperTest(unittest.TestCase):
|
|||
True,
|
||||
False,
|
||||
{},
|
||||
fusion_options,
|
||||
)
|
||||
model_path = base_path + model_name.replace("-", "_") + "_1.onnx"
|
||||
import onnx
|
||||
|
||||
return onnx.load_model(model_path)
|
||||
|
||||
# TODO: use a static lightweight model for test
|
||||
|
|
|
|||
Loading…
Reference in a new issue