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:
KeDengMS 2020-10-30 19:28:47 -07:00 committed by GitHub
parent 7a80a4b526
commit 32bf6390ad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 7 deletions

View file

@ -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")

View file

@ -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)