diff --git a/dockerfiles/Dockerfile.nuphar b/dockerfiles/Dockerfile.nuphar index f2d30ef36c..65be862957 100644 --- a/dockerfiles/Dockerfile.nuphar +++ b/dockerfiles/Dockerfile.nuphar @@ -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 diff --git a/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py b/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py index 48becfede3..c4c212a6ba 100644 --- a/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py +++ b/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py @@ -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!') diff --git a/onnxruntime/test/python/onnxruntime_test_python_nuphar.py b/onnxruntime/test/python/onnxruntime_test_python_nuphar.py index a72050d70f..c33226e632 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_nuphar.py +++ b/onnxruntime/test/python/onnxruntime_test_python_nuphar.py @@ -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()