mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Symbolic shape inference improvements: (#2460)
* Symbolic shape inference improvements: - add a mode to guess unknown ops' output rank - add support for GatherND - add support for If - fix a bug in get_int_values when then tensor rank > 1D, by treating it as no sympy data - add symbol to literal merge when ONNX silently merges dims - fix a bug in Concat when input dim is 0 - fix a bug in ConstantOfShape that computed dim is not updated - add support for dynamic shape in ConstantOfShape - fix a bug in Loop output shape that loop iterator dim is not inserted at dim 0 - add support for dynamic padding in Pad - add support for dynamic shape in Reshape - add support for Resize with opset > 10, by treating output dims as dynamic - fix a bug in Slice when starts/ends are dynamic - restrict input model to opset 7 and above - make output model optional to avoid disk write when testing Run model tests for symbolic shape inference Reduce 2GB docker image size of nuphar
This commit is contained in:
KeDengMS
GitHub
parent
a01816335d
commit
e0b51e4d95
3 changed files with 262 additions and 106 deletions
|
|
@ -22,8 +22,9 @@ RUN /onnxruntime/tools/ci_build/github/linux/docker/scripts/install_ubuntu.sh -p
|
|||
WORKDIR /
|
||||
|
||||
RUN mkdir -p /onnxruntime/build && \
|
||||
pip3 install sympy packaging && \
|
||||
python3 /onnxruntime/tools/ci_build/build.py --build_dir /onnxruntime/build --config Release --build_shared_lib --skip_submodule_sync --build_wheel --parallel --use_nuphar --use_mklml --use_tvm --use_llvm
|
||||
|
||||
RUN pip3 install /onnxruntime/build/Release/dist/onnxruntime_nuphar-*.whl && \
|
||||
pip3 install sympy packaging cpufeature jupyter && \
|
||||
python3 /onnxruntime/tools/ci_build/build.py --build_dir /onnxruntime/build --config Release --build_shared_lib --skip_submodule_sync --build_wheel --parallel --use_nuphar --use_mklml --use_tvm --use_llvm && \
|
||||
rm -rf /tmp/* && \
|
||||
rm -rf /home/root/* && \
|
||||
pip3 install /onnxruntime/build/Release/dist/onnxruntime_nuphar-*.whl && \
|
||||
rm -rf /onnxruntime
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from onnx import helper, numpy_helper, shape_inference
|
|||
import sympy
|
||||
|
||||
from packaging import version
|
||||
assert version.parse(onnx.__version__) >= version.parse("1.5.0") # need at least opset 10 for MatMulInteger shape inference
|
||||
assert version.parse(onnx.__version__) >= version.parse("1.5.0")
|
||||
|
||||
def get_attribute(node, attr_name, default_value=None):
|
||||
found = [attr for attr in node.attribute if attr.name == attr_name]
|
||||
|
|
@ -18,8 +18,11 @@ def get_attribute(node, attr_name, default_value=None):
|
|||
return helper.get_attribute_value(found[0])
|
||||
return default_value
|
||||
|
||||
def get_dim_from_type_proto(dim):
|
||||
return getattr(dim, dim.WhichOneof('value')) if type(dim.WhichOneof('value')) == str else None
|
||||
|
||||
def get_shape_from_type_proto(type_proto):
|
||||
return [getattr(i, i.WhichOneof('value')) if type(i.WhichOneof('value')) == str else None for i in type_proto.tensor_type.shape.dim]
|
||||
return [get_dim_from_type_proto(d) for d in type_proto.tensor_type.shape.dim]
|
||||
|
||||
def get_shape_from_sympy_shape(sympy_shape):
|
||||
return [None if i is None else (int(i) if is_literal(i) else str(i)) for i in sympy_shape]
|
||||
|
|
@ -48,11 +51,13 @@ def as_scalar(x):
|
|||
else:
|
||||
return x
|
||||
|
||||
def as_list(x):
|
||||
def as_list(x, keep_none):
|
||||
if type(x) == list:
|
||||
return x
|
||||
elif type(x) == np.ndarray:
|
||||
return list(x)
|
||||
elif keep_none and x is None:
|
||||
return None
|
||||
else:
|
||||
return [x]
|
||||
|
||||
|
|
@ -66,7 +71,7 @@ def sympy_reduce_product(x):
|
|||
return value
|
||||
|
||||
class SymbolicShapeInference:
|
||||
def __init__(self, int_max, auto_merge, verbose):
|
||||
def __init__(self, int_max, auto_merge, guess_output_rank, verbose):
|
||||
self.dispatcher_ = {
|
||||
'Add' : self._infer_binary_ops,
|
||||
'ArrayFeatureExtractor' : self._infer_ArrayFeatureExtractor,
|
||||
|
|
@ -82,6 +87,8 @@ class SymbolicShapeInference:
|
|||
'Expand' : self._infer_Expand,
|
||||
'Gather' : self._infer_Gather,
|
||||
'GatherElements' : self._infer_GatherElements,
|
||||
'GatherND' : self._infer_GatherND,
|
||||
'If' : self._infer_If,
|
||||
'Loop' : self._infer_Loop,
|
||||
'MatMul' : self._infer_MatMul,
|
||||
'MatMulInteger16' : self._infer_MatMulInteger,
|
||||
|
|
@ -114,6 +121,7 @@ class SymbolicShapeInference:
|
|||
self.suggested_merge_ = {}
|
||||
self.symbolic_dims_ = {}
|
||||
self.auto_merge_ = auto_merge
|
||||
self.guess_output_rank_ = guess_output_rank
|
||||
self.verbose_ = verbose
|
||||
self.int_max_ = int_max
|
||||
|
||||
|
|
@ -136,11 +144,15 @@ class SymbolicShapeInference:
|
|||
if type(self.symbolic_dims_[s]) == sympy.Symbol:
|
||||
map_to = s
|
||||
break
|
||||
# when nothing to map to, use the first one
|
||||
# when nothing to map to, use the shorter one
|
||||
if map_to is None:
|
||||
if self.verbose_ > 0:
|
||||
print('Potential unsafe merge between symbolic expressions: ({})'.format(','.join(symbols)))
|
||||
map_to = symbols.pop() # force merge when unable to determine
|
||||
symbols_list = list(symbols)
|
||||
lens = [len(s) for s in symbols_list]
|
||||
map_to = symbols_list[lens.index(min(lens))]
|
||||
symbols.remove(map_to)
|
||||
|
||||
for s in symbols:
|
||||
if s == map_to:
|
||||
continue
|
||||
|
|
@ -313,7 +325,7 @@ class SymbolicShapeInference:
|
|||
vi.CopyFrom(self.tmp_mp_.graph.output[i_o])
|
||||
self.known_vi_[o] = vi
|
||||
|
||||
def _onnx_infer_subgraph(self, node, subgraph):
|
||||
def _onnx_infer_subgraph(self, node, subgraph, use_node_input=True):
|
||||
if self.verbose_ > 2:
|
||||
print('Inferencing subgraph of node {} with output({}...): {}'.format(node.name, node.output[0], node.op_type))
|
||||
# node inputs are not passed directly to the subgraph
|
||||
|
|
@ -321,9 +333,7 @@ class SymbolicShapeInference:
|
|||
# for example, with Scan/Loop, subgraph input shape would be trimmed from node input shape
|
||||
# besides, inputs in subgraph could shadow implicit inputs
|
||||
subgraph_inputs = set([i.name for i in list(subgraph.initializer) + list(subgraph.input)])
|
||||
subgraph_implicit_input = set()
|
||||
for sn in subgraph.node:
|
||||
subgraph_implicit_input.update([i for i in sn.input if i in self.known_vi_ and i not in subgraph_inputs])
|
||||
subgraph_implicit_input = set([name for name in self.known_vi_.keys() if not name in subgraph_inputs])
|
||||
tmp_graph = helper.make_graph(list(subgraph.node),
|
||||
'tmp',
|
||||
list(subgraph.input) + [self.known_vi_[i] for i in subgraph_implicit_input],
|
||||
|
|
@ -332,21 +342,32 @@ class SymbolicShapeInference:
|
|||
tmp_graph.initializer.extend(subgraph.initializer)
|
||||
self.tmp_mp_.graph.CopyFrom(tmp_graph)
|
||||
|
||||
symbolic_shape_inference = SymbolicShapeInference(self.int_max_, self.auto_merge_, self.verbose_)
|
||||
symbolic_shape_inference = SymbolicShapeInference(self.int_max_, self.auto_merge_, self.guess_output_rank_, self.verbose_)
|
||||
all_shapes_inferred = False
|
||||
symbolic_shape_inference._preprocess(self.tmp_mp_)
|
||||
symbolic_shape_inference.suggested_merge_ = self.suggested_merge_.copy()
|
||||
while symbolic_shape_inference.run_:
|
||||
all_shapes_inferred = symbolic_shape_inference._infer_impl(self.tmp_mp_)
|
||||
all_shapes_inferred = symbolic_shape_inference._infer_impl(self.tmp_mp_, self.sympy_data_.copy())
|
||||
symbolic_shape_inference._update_output_from_vi()
|
||||
subgraph.ClearField('input')
|
||||
subgraph.input.extend(symbolic_shape_inference.out_mp_.graph.input[:len(node.input)])
|
||||
if use_node_input:
|
||||
# if subgraph uses node input, it needs to update to merged dims
|
||||
subgraph.ClearField('input')
|
||||
subgraph.input.extend(symbolic_shape_inference.out_mp_.graph.input[:len(node.input)])
|
||||
subgraph.ClearField('output')
|
||||
subgraph.output.extend(symbolic_shape_inference.out_mp_.graph.output)
|
||||
subgraph.ClearField('value_info')
|
||||
subgraph.value_info.extend(symbolic_shape_inference.out_mp_.graph.value_info)
|
||||
subgraph.ClearField('node')
|
||||
subgraph.node.extend(symbolic_shape_inference.out_mp_.graph.node)
|
||||
# for new symbolic dims from subgraph output, add to main graph symbolic dims
|
||||
subgraph_shapes = [get_shape_from_type_proto(o.type) for o in symbolic_shape_inference.out_mp_.graph.output]
|
||||
subgraph_new_symbolic_dims = set([d for s in subgraph_shapes if s for d in s if type(d) == str and not d in self.symbolic_dims_])
|
||||
self.symbolic_dims_.update({d:symbolic_shape_inference.symbolic_dims_[d] for d in subgraph_new_symbolic_dims})
|
||||
new_dims = {}
|
||||
for d in subgraph_new_symbolic_dims:
|
||||
assert d in symbolic_shape_inference.symbolic_dims_
|
||||
new_dims[d] = symbolic_shape_inference.symbolic_dims_[d]
|
||||
self.symbolic_dims_.update(new_dims)
|
||||
return symbolic_shape_inference
|
||||
|
||||
def _get_int_values(self, node, broadcast=False):
|
||||
values = [self._try_get_value(node, i) for i in range(len(node.input))]
|
||||
|
|
@ -355,8 +376,9 @@ class SymbolicShapeInference:
|
|||
for i,v in enumerate(values):
|
||||
if type(v) != np.ndarray:
|
||||
continue
|
||||
assert len(v.shape) <= 1
|
||||
if len(v.shape) == 0:
|
||||
if len(v.shape) > 1:
|
||||
new_v = None # ignore value for rank > 1
|
||||
elif len(v.shape) == 0:
|
||||
new_v = int(np.asscalar(v))
|
||||
else:
|
||||
assert len(v.shape) == 1
|
||||
|
|
@ -399,8 +421,8 @@ class SymbolicShapeInference:
|
|||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
self._get_shape(node, 0)))
|
||||
|
||||
def _new_symbolic_dim_from_output(self, node, out_idx=0, dim=0):
|
||||
new_dim = '{}{}_o{}_d{}'.format(node.op_type, list(self.out_mp_.graph.node).index(node), out_idx, dim)
|
||||
def _new_symbolic_dim(self, prefix, dim):
|
||||
new_dim = '{}_d{}'.format(prefix, dim)
|
||||
if new_dim in self.suggested_merge_:
|
||||
v = self.suggested_merge_[new_dim]
|
||||
new_dim = sympy.Integer(int(v)) if is_literal(v) else v
|
||||
|
|
@ -408,6 +430,12 @@ class SymbolicShapeInference:
|
|||
self.symbolic_dims_[new_dim] = sympy.Symbol(new_dim, integer=True)
|
||||
return new_dim
|
||||
|
||||
def _new_symbolic_dim_from_output(self, node, out_idx=0, dim=0):
|
||||
return self._new_symbolic_dim('{}{}_o{}_'.format(node.op_type, list(self.out_mp_.graph.node).index(node), out_idx), dim)
|
||||
|
||||
def _new_symbolic_shape(self, rank, node, out_idx=0):
|
||||
return [self._new_symbolic_dim_from_output(node, out_idx, i) for i in range(rank)]
|
||||
|
||||
def _compute_conv_pool_shape(self, node):
|
||||
sympy_shape = self._get_sympy_shape(node, 0)
|
||||
if len(node.input) > 1:
|
||||
|
|
@ -464,7 +492,13 @@ class SymbolicShapeInference:
|
|||
strided_kernel_positions = (effective_input_size - effective_kernel_shape[i]) // strides[i]
|
||||
sympy_shape[-rank + i] = strided_kernel_positions + 1
|
||||
return sympy_shape
|
||||
|
||||
|
||||
def _check_merged_dims(self, dims, allow_broadcast=True):
|
||||
if allow_broadcast:
|
||||
dims = [d for d in dims if not(is_literal(d) and int(d) <= 1)]
|
||||
if not all([d == dims[0] for d in dims]):
|
||||
self._add_suggested_merge(dims, apply=True)
|
||||
|
||||
def _compute_matmul_shape(self, node, output_dtype=None):
|
||||
lhs_shape = self._get_shape(node, 0)
|
||||
rhs_shape = self._get_shape(node, 1)
|
||||
|
|
@ -485,10 +519,8 @@ class SymbolicShapeInference:
|
|||
lhs_reduce_dim = -1
|
||||
rhs_reduce_dim = -2
|
||||
new_shape = self._broadcast_shapes(lhs_shape[:-2], rhs_shape[:-2]) + [lhs_shape[-2]] + [rhs_shape[-1]]
|
||||
# record inconsistent reduce dim as suggested merge
|
||||
if lhs_shape[lhs_reduce_dim] != rhs_shape[rhs_reduce_dim]:
|
||||
merge_dims = [lhs_shape[lhs_reduce_dim], rhs_shape[rhs_reduce_dim]]
|
||||
self._add_suggested_merge(merge_dims, apply=True)
|
||||
# merge reduce dim
|
||||
self._check_merged_dims([lhs_shape[lhs_reduce_dim], rhs_shape[rhs_reduce_dim]], allow_broadcast=False)
|
||||
if output_dtype is None:
|
||||
# infer output_dtype from input type when not specified
|
||||
output_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type
|
||||
|
|
@ -557,13 +589,15 @@ class SymbolicShapeInference:
|
|||
sympy_shape = self._get_sympy_shape(node, 0)
|
||||
axis = handle_negative_axis(get_attribute(node, 'axis'), len(sympy_shape))
|
||||
for i_idx in range(1, len(node.input)):
|
||||
sympy_shape[axis] = sympy_shape[axis] + self._get_sympy_shape(node, i_idx)[axis]
|
||||
input_shape = self._get_sympy_shape(node, i_idx)
|
||||
if input_shape:
|
||||
sympy_shape[axis] = sympy_shape[axis] + input_shape[axis]
|
||||
self._update_computed_dims(sympy_shape)
|
||||
# merge symbolic dims for non-concat axes
|
||||
for d in range(len(sympy_shape)):
|
||||
if d == axis:
|
||||
continue
|
||||
dims = [self._get_shape(node, i_idx)[d] for i_idx in range(len(node.input))]
|
||||
dims = [self._get_shape(node, i_idx)[d] for i_idx in range(len(node.input)) if self._get_shape(node, i_idx)]
|
||||
if all([d == dims[0] for d in dims]):
|
||||
continue
|
||||
merged = self._merge_symbols(dims)
|
||||
|
|
@ -582,13 +616,19 @@ class SymbolicShapeInference:
|
|||
|
||||
def _infer_ConstantOfShape(self, node):
|
||||
sympy_shape = self._get_int_values(node)[0]
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
if sympy_shape is not None:
|
||||
self._update_computed_dims(sympy_shape)
|
||||
if type(sympy_shape) != list:
|
||||
sympy_shape = [sympy_shape]
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
vi.type.tensor_type.elem_type,
|
||||
get_shape_from_sympy_shape(sympy_shape)))
|
||||
else:
|
||||
# create new dynamic shape
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
vi.type.tensor_type.elem_type,
|
||||
self._new_symbolic_shape(self._get_shape_rank(node,0), node)))
|
||||
|
||||
def _infer_Expand(self, node):
|
||||
expand_to_shape = self._try_get_value(node, 1)
|
||||
|
|
@ -627,6 +667,43 @@ class SymbolicShapeInference:
|
|||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
indices_shape))
|
||||
|
||||
def _infer_GatherND(self, node):
|
||||
data_shape = self._get_shape(node, 0)
|
||||
data_rank = len(data_shape)
|
||||
indices_shape = self._get_shape(node, 1)
|
||||
indices_rank = len(indices_shape)
|
||||
last_index_dimension = indices_shape[-1]
|
||||
assert is_literal(last_index_dimension) and last_index_dimension <= data_rank
|
||||
new_shape = indices_shape[:-1] + data_shape[last_index_dimension:]
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
new_shape))
|
||||
|
||||
def _infer_If(self, node):
|
||||
# special case for constant condition, in case there are mismatching shape from the non-executed branch
|
||||
subgraphs = [get_attribute(node, 'then_branch'), get_attribute(node, 'else_branch')]
|
||||
cond = self._try_get_value(node, 0)
|
||||
if cond is not None:
|
||||
if cond > 0:
|
||||
subgraphs[1].CopyFrom(subgraphs[0])
|
||||
else:
|
||||
subgraphs[0].CopyFrom(subgraphs[1])
|
||||
|
||||
for i_sub, subgraph in enumerate(subgraphs):
|
||||
subgraph_infer = self._onnx_infer_subgraph(node, subgraph, use_node_input=False)
|
||||
for i_out in range(len(node.output)):
|
||||
vi = self.known_vi_[node.output[i_out]]
|
||||
if i_sub == 0:
|
||||
vi.CopyFrom(subgraph.output[i_out])
|
||||
vi.name = node.output[i_out]
|
||||
else:
|
||||
assert all([d1 == d2 for d1,d2 in zip(vi.type.tensor_type.shape.dim, subgraph.output[i_out].type.tensor_type.shape.dim)])
|
||||
# pass on sympy data from subgraph, if cond is constant
|
||||
if cond is not None and i_sub == (0 if cond > 0 else 1):
|
||||
if subgraph.output[i_out].name in subgraph_infer.sympy_data_:
|
||||
self.sympy_data_[vi.name] = subgraph_infer.sympy_data_[subgraph.output[i_out].name]
|
||||
|
||||
def _infer_Loop(self, node):
|
||||
subgraph = get_attribute(node, 'body')
|
||||
assert len(subgraph.input) == len(node.input)
|
||||
|
|
@ -642,11 +719,11 @@ class SymbolicShapeInference:
|
|||
vi = self.known_vi_[node.output[i]]
|
||||
vi.CopyFrom(subgraph.output[i + 1]) # first subgraph output is condition, not in node output
|
||||
if i >= num_loop_carried:
|
||||
subgraph_vi_dim = subgraph.output[i + 1].type.tensor_type.shape.dim
|
||||
vi.type.tensor_type.shape.ClearField('dim')
|
||||
vi_dim = vi.type.tensor_type.shape.dim
|
||||
if len(vi_dim) > 0:
|
||||
vi_dim[0].dim_param = loop_iter_dim
|
||||
else:
|
||||
vi_dim.add().dim_param = loop_iter_dim
|
||||
vi_dim.add().dim_param = loop_iter_dim
|
||||
vi_dim.extend(list(subgraph_vi_dim))
|
||||
vi.name = node.output[i]
|
||||
|
||||
def _infer_MatMul(self, node):
|
||||
|
|
@ -679,16 +756,20 @@ class SymbolicShapeInference:
|
|||
if get_opset(self.out_mp_) <= 10:
|
||||
pads = get_attribute(node, 'pads')
|
||||
else:
|
||||
pads = self._get_value(node, 1)
|
||||
pads = self._try_get_value(node, 1)
|
||||
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
output_shape = get_shape_from_type_proto(vi.type)
|
||||
if len(output_shape) == 0 or None in output_shape:
|
||||
sympy_shape = self._get_sympy_shape(node, 0)
|
||||
rank = len(sympy_shape)
|
||||
assert len(pads) == 2*rank
|
||||
new_shape = [d + pad_up + pad_down for d, pad_up, pad_down in zip(sympy_shape, pads[:rank], pads[rank:])]
|
||||
self._update_computed_dims(new_shape)
|
||||
if pads is not None:
|
||||
assert len(pads) == 2*rank
|
||||
new_shape = [d + pad_up + pad_down for d, pad_up, pad_down in zip(sympy_shape, pads[:rank], pads[rank:])]
|
||||
self._update_computed_dims(new_shape)
|
||||
else:
|
||||
# dynamic pads, create new symbolic dimensions
|
||||
new_shape = self._new_symbolic_shape(rank, node)
|
||||
output_tp = self.known_vi_[node.input[0]].type.tensor_type.elem_type
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_tp, get_shape_from_sympy_shape(new_shape)))
|
||||
|
||||
|
|
@ -723,49 +804,65 @@ class SymbolicShapeInference:
|
|||
self.sympy_data_[node.output[0]] = sympy_reduce_product(data)
|
||||
|
||||
def _infer_Reshape(self, node):
|
||||
shape_value = self._get_value(node, 1)
|
||||
input_shape = self._get_shape(node, 0)
|
||||
input_sympy_shape = self._get_sympy_shape(node, 0)
|
||||
total = int(1)
|
||||
for d in input_sympy_shape:
|
||||
total = total * d
|
||||
new_sympy_shape = []
|
||||
deferred_dim_idx = -1
|
||||
non_deferred_size = int(1)
|
||||
for i, d in enumerate(shape_value):
|
||||
if type(d) == sympy.Symbol:
|
||||
new_sympy_shape.append(d)
|
||||
elif d == 0:
|
||||
new_sympy_shape.append(input_sympy_shape[i])
|
||||
non_deferred_size = non_deferred_size * input_sympy_shape[i]
|
||||
else:
|
||||
new_sympy_shape.append(d)
|
||||
if d == -1:
|
||||
deferred_dim_idx = i
|
||||
elif d != 0:
|
||||
non_deferred_size = non_deferred_size * d
|
||||
|
||||
assert new_sympy_shape.count(-1) < 2
|
||||
if -1 in new_sympy_shape:
|
||||
new_dim = total // non_deferred_size
|
||||
new_sympy_shape[deferred_dim_idx] = new_dim
|
||||
self._update_computed_dims(new_sympy_shape)
|
||||
|
||||
shape_value = self._try_get_value(node, 1)
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
vi.type.tensor_type.elem_type,
|
||||
get_shape_from_sympy_shape(new_sympy_shape)))
|
||||
if shape_value is None:
|
||||
shape_shape = self._get_shape(node, 1)
|
||||
assert len(shape_shape) == 1
|
||||
shape_rank = shape_shape[0]
|
||||
assert is_literal(shape_rank)
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
vi.type.tensor_type.elem_type,
|
||||
self._new_symbolic_shape(shape_rank, node)))
|
||||
else:
|
||||
input_shape = self._get_shape(node, 0)
|
||||
input_sympy_shape = self._get_sympy_shape(node, 0)
|
||||
total = int(1)
|
||||
for d in input_sympy_shape:
|
||||
total = total * d
|
||||
new_sympy_shape = []
|
||||
deferred_dim_idx = -1
|
||||
non_deferred_size = int(1)
|
||||
for i, d in enumerate(shape_value):
|
||||
if type(d) == sympy.Symbol:
|
||||
new_sympy_shape.append(d)
|
||||
elif d == 0:
|
||||
new_sympy_shape.append(input_sympy_shape[i])
|
||||
non_deferred_size = non_deferred_size * input_sympy_shape[i]
|
||||
else:
|
||||
new_sympy_shape.append(d)
|
||||
if d == -1:
|
||||
deferred_dim_idx = i
|
||||
elif d != 0:
|
||||
non_deferred_size = non_deferred_size * d
|
||||
|
||||
assert new_sympy_shape.count(-1) < 2
|
||||
if -1 in new_sympy_shape:
|
||||
new_dim = total // non_deferred_size
|
||||
new_sympy_shape[deferred_dim_idx] = new_dim
|
||||
self._update_computed_dims(new_sympy_shape)
|
||||
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
vi.type.tensor_type.elem_type,
|
||||
get_shape_from_sympy_shape(new_sympy_shape)))
|
||||
|
||||
self._pass_on_sympy_data(node)
|
||||
|
||||
def _infer_Resize(self, node):
|
||||
assert get_opset(self.out_mp_) <= 10 # only support opset 10 Resize for now
|
||||
scales = self._try_get_value(node, 1)
|
||||
if scales is not None:
|
||||
input_sympy_shape = self._get_sympy_shape(node, 0)
|
||||
new_sympy_shape = [sympy.simplify(sympy.floor(d*s)) for d,s in zip(input_sympy_shape, scales)]
|
||||
self._update_computed_dims(new_sympy_shape)
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[0]].type.tensor_type.elem_type, get_shape_from_sympy_shape(new_sympy_shape)))
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
if get_opset(self.out_mp_) <= 10: # only support opset 10 Resize for now
|
||||
scales = self._try_get_value(node, 1)
|
||||
if scales is not None:
|
||||
input_sympy_shape = self._get_sympy_shape(node, 0)
|
||||
new_sympy_shape = [sympy.simplify(sympy.floor(d*s)) for d,s in zip(input_sympy_shape, scales)]
|
||||
self._update_computed_dims(new_sympy_shape)
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
get_shape_from_sympy_shape(new_sympy_shape)))
|
||||
else:
|
||||
vi.CopyFrom(helper.make_tensor_value_info(node.output[0],
|
||||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
self._new_symbolic_shape(self._get_shape_rank(node, 0), node)))
|
||||
|
||||
def _infer_Scan(self, node):
|
||||
subgraph = get_attribute(node, 'body')
|
||||
|
|
@ -817,16 +914,16 @@ class SymbolicShapeInference:
|
|||
ends = get_attribute(node, 'ends')
|
||||
steps = [1]*len(axes)
|
||||
else:
|
||||
starts = as_list(self._try_get_value(node, 1))
|
||||
ends = as_list(self._try_get_value(node, 2))
|
||||
starts = as_list(self._try_get_value(node, 1), keep_none=True)
|
||||
ends = as_list(self._try_get_value(node, 2), keep_none=True)
|
||||
axes = self._try_get_value(node, 3)
|
||||
steps = self._try_get_value(node, 4)
|
||||
if axes is None and not (starts is None and ends is None):
|
||||
axes = list(range(0, len(starts if starts is not None else ends)))
|
||||
if steps is None and not (starts is None and ends is None):
|
||||
steps = [1]*len(starts if starts is not None else ends)
|
||||
axes = as_list(axes)
|
||||
steps = as_list(steps)
|
||||
axes = as_list(axes, keep_none=True)
|
||||
steps = as_list(steps, keep_none=True)
|
||||
|
||||
new_sympy_shape = self._get_sympy_shape(node, 0)
|
||||
if starts is None or ends is None:
|
||||
|
|
@ -844,7 +941,7 @@ class SymbolicShapeInference:
|
|||
if e >= self.int_max_:
|
||||
e = new_sympy_shape[i]
|
||||
elif e <= -self.int_max_:
|
||||
e = 0 if step > 0 else -1
|
||||
e = 0 if s > 0 else -1
|
||||
elif is_literal(new_sympy_shape[i]):
|
||||
if e < 0:
|
||||
e = e + new_sympy_shape[i]
|
||||
|
|
@ -963,13 +1060,19 @@ class SymbolicShapeInference:
|
|||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(new_vi)
|
||||
|
||||
def _infer_impl(self, in_mp):
|
||||
self.sympy_data_ = {}
|
||||
def _infer_impl(self, in_mp, start_sympy_data={}):
|
||||
self.sympy_data_ = start_sympy_data
|
||||
self.out_mp_.graph.ClearField('value_info')
|
||||
self._apply_suggested_merge(graph_input_only=True)
|
||||
input_symbols = set()
|
||||
for i in self.out_mp_.graph.input:
|
||||
input_dims = i.type.tensor_type.shape.dim
|
||||
for i_dim in range(len(input_dims)):
|
||||
if get_dim_from_type_proto(input_dims[i_dim]) is None:
|
||||
# some models use None for symbolic dim in input, replace it with a string
|
||||
input_dims[i_dim].dim_param = self._new_symbolic_dim(i.name, i_dim)
|
||||
input_symbols.update([d for d in get_shape_from_type_proto(i.type) if type(d) == str])
|
||||
|
||||
for s in input_symbols:
|
||||
if s in self.suggested_merge_:
|
||||
s_merge = self.suggested_merge_[s]
|
||||
|
|
@ -993,18 +1096,34 @@ class SymbolicShapeInference:
|
|||
|
||||
if self.verbose_ > 2:
|
||||
print(node.op_type + ': ' + node.name)
|
||||
for i, name in enumerate(node.input):
|
||||
print(' Input {}: {} {}'.format(i, name, 'initializer' if name in self.initializers_ else ''))
|
||||
|
||||
# onnx automatically merge dims with value, i.e. Mul(['aaa', 'bbb'], [1000, 1]) -> [1000, 'bbb']
|
||||
# symbolic shape inference needs to apply merge of 'aaa' -> 1000 in this case
|
||||
if node.op_type in ['Add', 'Sub', 'Mul', 'Div', 'MatMul', 'MatMulInteger', 'MatMulInteger16', 'Where', 'Sum']:
|
||||
vi = self.known_vi_[node.output[0]]
|
||||
out_rank = len(get_shape_from_type_proto(vi.type))
|
||||
in_shapes = [self._get_shape(node, i) for i in range(len(node.input))]
|
||||
for d in range(out_rank - (2 if node.op_type in ['MatMul', 'MatMulInteger', 'MatMulInteger16'] else 0)):
|
||||
in_dims = [s[len(s) - out_rank + d] for s in in_shapes if len(s) + d >= out_rank]
|
||||
if len(in_dims) > 1:
|
||||
self._check_merged_dims(in_dims, allow_broadcast=True)
|
||||
|
||||
for i_o in range(len(node.output)):
|
||||
out_type = self.known_vi_[node.output[i_o]].type
|
||||
vi = self.known_vi_[node.output[i_o]]
|
||||
out_type = vi.type
|
||||
out_type_kind = out_type.WhichOneof('value')
|
||||
# only TensorProto and SparseTensorProto have shape
|
||||
if out_type_kind != 'tensor_type' and out_type_kind != 'sparse_tensor_type':
|
||||
continue
|
||||
out_shape = get_shape_from_type_proto(self.known_vi_[node.output[i_o]].type)
|
||||
out_shape = get_shape_from_type_proto(vi.type)
|
||||
out_type_undefined = out_type.tensor_type.elem_type == onnx.TensorProto.UNDEFINED
|
||||
if self.verbose_ > 2:
|
||||
print(' {}: {} {}'.format(node.output[i_o], str(out_shape), self.known_vi_[node.output[i_o]].type.tensor_type.elem_type))
|
||||
print(' {}: {} {}'.format(node.output[i_o], str(out_shape), vi.type.tensor_type.elem_type))
|
||||
if node.output[i_o] in self.sympy_data_:
|
||||
print(' Sympy Data: ' + str(self.sympy_data_[node.output[i_o]]))
|
||||
|
||||
if None in out_shape or out_type_undefined:
|
||||
if self.auto_merge_:
|
||||
if node.op_type in ['Add', 'Sub', 'Mul', 'Div', 'MatMul', 'MatMulInteger', 'MatMulInteger16', 'Concat', 'Where', 'Sum']:
|
||||
|
|
@ -1032,9 +1151,34 @@ class SymbolicShapeInference:
|
|||
else:
|
||||
self.run_ = False
|
||||
|
||||
# create new dynamic dims for ops not handled by symbolic shape inference
|
||||
if self.run_ == False and not node.op_type in self.dispatcher_:
|
||||
is_unknown_op = (out_type_undefined and len(out_shape) == 0)
|
||||
if is_unknown_op:
|
||||
# unknown op to ONNX, maybe from higher opset or other domain
|
||||
# only guess the output rank from input 0 when using guess_output_rank option
|
||||
out_rank = self._get_shape_rank(node, 0) if self.guess_output_rank_ else -1
|
||||
else:
|
||||
# valid ONNX op, but not handled by symbolic shape inference, just assign dynamic shape
|
||||
out_rank = len(out_shape)
|
||||
|
||||
if out_rank >= 0:
|
||||
new_shape = self._new_symbolic_shape(out_rank, node, i_o)
|
||||
vi.CopyFrom(helper.make_tensor_value_info(vi.name,
|
||||
self.known_vi_[node.input[0]].type.tensor_type.elem_type,
|
||||
new_shape))
|
||||
|
||||
if self.verbose_ > 0:
|
||||
if is_unknown_op:
|
||||
print("Possible unknown op: {} node: {}, guessing {} shape".format(node.op_type, node.name, vi.name))
|
||||
if self.verbose_ > 2:
|
||||
print(' {}: {} {}'.format(node.output[i_o], str(new_shape), vi.type.tensor_type.elem_type))
|
||||
|
||||
self.run_ = True
|
||||
continue # continue the inference after guess, no need to stop as no merge is needed
|
||||
|
||||
if self.verbose_ > 0 or not self.auto_merge_ or out_type_undefined:
|
||||
print('Stopping at incomplete shape inference at ' + node.op_type + ': ' + node.name)
|
||||
print(node)
|
||||
print('node inputs:')
|
||||
for i in node.input:
|
||||
print(self.known_vi_[i])
|
||||
|
|
@ -1054,31 +1198,37 @@ class SymbolicShapeInference:
|
|||
output.CopyFrom(self.known_vi_[output.name])
|
||||
|
||||
@staticmethod
|
||||
def infer_shapes(input_model, output_model, int_max=2**31 - 1, auto_merge=False, verbose=0):
|
||||
def infer_shapes(input_model, output_model, int_max=2**31 - 1, auto_merge=False, guess_output_rank=False, verbose=0):
|
||||
in_mp = onnx.load(input_model)
|
||||
symbolic_shape_inference = SymbolicShapeInference(int_max, auto_merge, verbose)
|
||||
if get_opset(in_mp) < 7:
|
||||
print('Only support models of opset 7 and above.')
|
||||
return
|
||||
symbolic_shape_inference = SymbolicShapeInference(int_max, auto_merge, guess_output_rank, verbose)
|
||||
all_shapes_inferred = False
|
||||
symbolic_shape_inference._preprocess(in_mp)
|
||||
while symbolic_shape_inference.run_:
|
||||
all_shapes_inferred = symbolic_shape_inference._infer_impl(in_mp)
|
||||
symbolic_shape_inference._update_output_from_vi()
|
||||
onnx.save(symbolic_shape_inference.out_mp_, output_model)
|
||||
if output_model:
|
||||
onnx.save(symbolic_shape_inference.out_mp_, output_model)
|
||||
if not all_shapes_inferred:
|
||||
sys.exit(1)
|
||||
|
||||
def parse_arguments():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input', required=True, help='The input model file')
|
||||
parser.add_argument('--output', required=True, help='The input model file')
|
||||
parser.add_argument('--output', help='The input model file')
|
||||
parser.add_argument('--auto_merge', help='Automatically merge symbolic dims when confliction happens', action='store_true', default=False)
|
||||
parser.add_argument('--int_max', help='maximum value for integer to be treated as boundless for ops like slice', type=int, default=2**31 - 1)
|
||||
parser.add_argument('--guess_output_rank', help='guess output rank to be the same as input 0 for unknown ops', action='store_true', default=False)
|
||||
parser.add_argument('--verbose', help='Prints detailed logs of inference, 0: turn off, 1: warnings, 3: detailed', type=int, default=0)
|
||||
return parser.parse_args()
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_arguments()
|
||||
print('input model: ' + args.input)
|
||||
print('output model ' + args.output)
|
||||
if args.output:
|
||||
print('output model ' + args.output)
|
||||
print('Doing symbolic shape inference...')
|
||||
out_mp = SymbolicShapeInference.infer_shapes(args.input, args.output, args.int_max, args.auto_merge, args.verbose)
|
||||
out_mp = SymbolicShapeInference.infer_shapes(args.input, args.output, args.int_max, args.auto_merge, args.guess_output_rank, args.verbose)
|
||||
print('Done!')
|
||||
|
|
|
|||
|
|
@ -8,10 +8,11 @@ from onnx import numpy_helper
|
|||
import onnxruntime as onnxrt
|
||||
import os
|
||||
from onnxruntime.nuphar.rnn_benchmark import perf_test
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import sys
|
||||
import subprocess
|
||||
import tarfile
|
||||
from timeit import default_timer as timer
|
||||
import unittest
|
||||
import urllib.request
|
||||
|
||||
|
|
@ -47,13 +48,8 @@ class TestNuphar(unittest.TestCase):
|
|||
# test AOT on the quantized model
|
||||
cache_dir = os.path.join(cwd, 'nuphar_cache')
|
||||
if os.path.exists(cache_dir):
|
||||
for sub_dir in os.listdir(cache_dir):
|
||||
full_sub_dir = os.path.join(cache_dir, sub_dir)
|
||||
if os.path.isdir(full_sub_dir):
|
||||
for f in os.listdir(full_sub_dir):
|
||||
os.remove(os.path.join(full_sub_dir, f))
|
||||
else:
|
||||
os.makedirs(cache_dir)
|
||||
shutil.rmtree(cache_dir)
|
||||
os.makedirs(cache_dir)
|
||||
|
||||
# prepare feed
|
||||
feed = {}
|
||||
|
|
@ -113,9 +109,9 @@ class TestNuphar(unittest.TestCase):
|
|||
subprocess.run([onnx_test_runner, '-e', 'nuphar', '-n', 'download_sample_10', cwd], check=True, cwd=cwd)
|
||||
|
||||
# run onnxruntime_perf_test
|
||||
onnx_test_runner = os.path.join(cwd, 'onnxruntime_perf_test')
|
||||
subprocess.run([onnx_test_runner, '-e', 'nuphar', '-t', '20', bert_squad_model, '1.txt'], check=True, cwd=cwd)
|
||||
subprocess.run([onnx_test_runner, '-e', 'cpu', '-o', '99', '-t', '20', bert_squad_model, '1.txt'], check=True, cwd=cwd)
|
||||
onnxruntime_perf_test = os.path.join(cwd, 'onnxruntime_perf_test')
|
||||
subprocess.run([onnxruntime_perf_test, '-e', 'nuphar', '-t', '20', bert_squad_model, '1.txt'], check=True, cwd=cwd)
|
||||
subprocess.run([onnxruntime_perf_test, '-e', 'cpu', '-o', '99', '-t', '20', bert_squad_model, '1.txt'], check=True, cwd=cwd)
|
||||
|
||||
|
||||
def test_rnn_benchmark(self):
|
||||
|
|
@ -135,5 +131,14 @@ class TestNuphar(unittest.TestCase):
|
|||
min_duration_seconds=1)
|
||||
|
||||
|
||||
def test_symbolic_shape_infer(self):
|
||||
cwd = os.getcwd()
|
||||
test_model_dir = os.path.join(cwd, '..', 'models')
|
||||
for filename in Path(test_model_dir).rglob('*.onnx'):
|
||||
if filename.name.startswith('.'):
|
||||
continue # skip some bad model files
|
||||
subprocess.run([sys.executable, '-m', 'onnxruntime.nuphar.symbolic_shape_infer', '--input', str(filename), '--auto_merge', '--int_max=100000', '--guess_output_rank'], check=True, cwd=cwd)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
|||
Loading…
Reference in a new issue