mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Some fixes to symbolic shape inference (#5642)
* Some fixes to symbolic shape inference 1. Topological sort before iteration in graph 2. Fix a case in slice: start=100000, end=-100000, step=-1, dim=2 3. Fix Nuphar Gemm test's random seed 4. Slice opset 1 axes is optional
This commit is contained in:
parent
7a80a4b526
commit
32bf6390ad
2 changed files with 27 additions and 7 deletions
|
|
@ -343,7 +343,7 @@ class SymbolicShapeInference:
|
|||
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_, self.sympy_data_.copy())
|
||||
all_shapes_inferred = symbolic_shape_inference._infer_impl(self.sympy_data_.copy())
|
||||
symbolic_shape_inference._update_output_from_vi()
|
||||
if use_node_input:
|
||||
# if subgraph uses node input, it needs to update to merged dims
|
||||
|
|
@ -947,6 +947,8 @@ class SymbolicShapeInference:
|
|||
axes = get_attribute(node, 'axes')
|
||||
starts = get_attribute(node, 'starts')
|
||||
ends = get_attribute(node, 'ends')
|
||||
if not axes:
|
||||
axes = list(range(len(starts)))
|
||||
steps = [1]*len(axes)
|
||||
else:
|
||||
starts = as_list(self._try_get_value(node, 1), keep_none=True)
|
||||
|
|
@ -971,7 +973,6 @@ class SymbolicShapeInference:
|
|||
new_sympy_shape[i] = self._new_symbolic_dim_from_output(node,0,i)
|
||||
else:
|
||||
for i,s,e,t in zip(axes, starts, ends, steps):
|
||||
idx = handle_negative_axis(i, len(new_sympy_shape))
|
||||
if is_literal(e):
|
||||
if e >= self.int_max_:
|
||||
e = new_sympy_shape[i]
|
||||
|
|
@ -979,7 +980,7 @@ class SymbolicShapeInference:
|
|||
e = 0 if s > 0 else -1
|
||||
elif is_literal(new_sympy_shape[i]):
|
||||
if e < 0:
|
||||
e = e + new_sympy_shape[i]
|
||||
e = max(0, e + new_sympy_shape[i])
|
||||
e = min(e, new_sympy_shape[i])
|
||||
else:
|
||||
if e > 0:
|
||||
|
|
@ -999,8 +1000,10 @@ class SymbolicShapeInference:
|
|||
|
||||
if is_literal(s) and int(s) < 0:
|
||||
s = new_sympy_shape[i] + s
|
||||
if is_literal(new_sympy_shape[i]) and is_literal(s):
|
||||
s = max(0, min(s, new_sympy_shape[i]))
|
||||
|
||||
new_sympy_shape[idx] = (e - s + t + (-1 if t > 0 else 1)) // t
|
||||
new_sympy_shape[i] = (e - s + t + (-1 if t > 0 else 1)) // t
|
||||
|
||||
self._update_computed_dims(new_sympy_shape)
|
||||
|
||||
|
|
@ -1110,7 +1113,7 @@ class SymbolicShapeInference:
|
|||
vi = self.known_vi_[node.output[0]]
|
||||
vi.CopyFrom(new_vi)
|
||||
|
||||
def _infer_impl(self, in_mp, start_sympy_data=None):
|
||||
def _infer_impl(self, start_sympy_data=None):
|
||||
self.sympy_data_ = start_sympy_data or {}
|
||||
self.out_mp_.graph.ClearField('value_info')
|
||||
self._apply_suggested_merge(graph_input_only=True)
|
||||
|
|
@ -1138,7 +1141,23 @@ class SymbolicShapeInference:
|
|||
self.tmp_mp_.CopyFrom(self.out_mp_)
|
||||
self.tmp_mp_.graph.ClearField('initializer')
|
||||
|
||||
for node in self.out_mp_.graph.node:
|
||||
# topological sort nodes, note there might be dead nodes so we check if all graph outputs are reached to terminate
|
||||
sorted_nodes = []
|
||||
sorted_known_vi = set([i.name for i in list(self.out_mp_.graph.input) + list(self.out_mp_.graph.initializer)])
|
||||
if all([o.name in sorted_known_vi for o in self.out_mp_.graph.output]):
|
||||
# Loop/Scan will have all graph output in graph inputs, so don't do topological sort
|
||||
sorted_nodes = self.out_mp_.graph.node
|
||||
else:
|
||||
while not all([o.name in sorted_known_vi for o in self.out_mp_.graph.output]):
|
||||
old_sorted_nodes_len = len(sorted_nodes)
|
||||
for node in self.out_mp_.graph.node:
|
||||
if (node.output[0] not in sorted_known_vi ) and all([i in sorted_known_vi for i in node.input if i]):
|
||||
sorted_known_vi.update(node.output)
|
||||
sorted_nodes.append(node)
|
||||
if old_sorted_nodes_len == len(sorted_nodes) and not all([o.name in sorted_known_vi for o in self.out_mp_.graph.output]):
|
||||
raise Exception('Invalid model with cyclic graph')
|
||||
|
||||
for node in sorted_nodes:
|
||||
assert all([i in self.known_vi_ for i in node.input if i])
|
||||
self._onnx_infer_single_node(node)
|
||||
if node.op_type in self.dispatcher_:
|
||||
|
|
@ -1275,7 +1294,7 @@ class SymbolicShapeInference:
|
|||
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)
|
||||
all_shapes_inferred = symbolic_shape_inference._infer_impl()
|
||||
symbolic_shape_inference._update_output_from_vi()
|
||||
if not all_shapes_inferred:
|
||||
raise Exception("Incomplete symbolic shape inference")
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ def generate_gemm_inputs_initializers(graph, config, added_inputs_initializers={
|
|||
input_shape_b = ['seq'] + shape_b if extend else shape_b
|
||||
input_shape_c = ['seq'] + shape_c if extend else shape_c
|
||||
|
||||
np.random.seed(12345)
|
||||
a = np.random.ranf(shape_a).astype(np.float32)
|
||||
b = np.random.ranf(shape_b).astype(np.float32)
|
||||
c = np.random.ranf(shape_c).astype(np.float32) if config['withC'] else np.array(0)
|
||||
|
|
|
|||
Loading…
Reference in a new issue