diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 5950664d39..fe8b1951e0 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -548,7 +548,13 @@ class SymbolicShapeInference: self.symbolic_dims_.update(new_dims) return symbolic_shape_inference - def _get_int_values(self, node, broadcast=False): + def _get_int_or_float_values(self, node, broadcast=False, allow_float_values=False): + def int_or_float(value, allow_float_values): + # If casting into int has precision loss: keep float output + if allow_float_values and value % 1 != 0: + return value + return int(value) + values = [self._try_get_value(node, i) for i in range(len(node.input))] if all([v is not None for v in values]): # some shape compute is in floating point, cast to int for sympy @@ -558,10 +564,10 @@ class SymbolicShapeInference: if len(v.shape) > 1: new_v = None # ignore value for rank > 1 elif len(v.shape) == 0: - new_v = int(v.item()) + new_v = int_or_float(v.item(), allow_float_values) else: assert len(v.shape) == 1 - new_v = [int(vv) for vv in v] + new_v = [int_or_float(vv, allow_float_values) for vv in v] values[i] = new_v values_len = [len(v) if type(v) == list else 0 for v in values] max_len = max(values_len) @@ -581,7 +587,15 @@ class SymbolicShapeInference: def _compute_on_sympy_data(self, node, op_func): assert len(node.output) == 1 - values = self._get_int_values(node, broadcast=True) + + # Before mul & div operations + # cast inputs into interger might lose decimal part and reduce precision + # keep them as float, finish the operation, then cast the result into integer + if node.op_type in ["Mul", "Div"]: + values = self._get_int_or_float_values(node, broadcast=True, allow_float_values=True) + else: + values = self._get_int_or_float_values(node, broadcast=True) + if all([v is not None for v in values]): is_list = [type(v) == list for v in values] as_list = any(is_list) @@ -781,7 +795,9 @@ class SymbolicShapeInference: def _infer_symbolic_compute_ops(self, node): funcs = { "Add": lambda l: l[0] + l[1], # noqa: E741 - "Div": lambda l: l[0] // l[1], # integer div in sympy # noqa: E741 + "Div": lambda l: int(l[0] // l[1]) # noqa: E741 + if isinstance(l[0] // l[1], float) + else l[0] // l[1], # integer div in sympy "Equal": lambda l: l[0] == l[1], # noqa: E741 "Floor": lambda l: sympy.floor(l[0]), # noqa: E741 "Max": lambda l: l[1] # noqa: E741 @@ -790,7 +806,7 @@ class SymbolicShapeInference: "Min": lambda l: l[1] # noqa: E741 if is_literal(l[0]) and int(l[0]) > self.int_max_ else (l[0] if is_literal(l[1]) and int(l[1]) > self.int_max_ else sympy.Min(l[0], l[1])), - "Mul": lambda l: l[0] * l[1], # noqa: E741 + "Mul": lambda l: int(l[0] * l[1]) if isinstance(l[0] * l[1], float) else l[0] * l[1], # noqa: E741 "Sub": lambda l: l[0] - l[1], # noqa: E741 "Where": lambda l: l[1] if l[0] else l[2], # noqa: E741 "Neg": lambda l: -l[0], # noqa: E741 @@ -832,7 +848,7 @@ class SymbolicShapeInference: def _infer_Concat(self, node): # noqa: N802 if any([i in self.sympy_data_ or i in self.initializers_ for i in node.input]): - values = self._get_int_values(node) + values = self._get_int_or_float_values(node) if all([v is not None for v in values]): assert get_attribute(node, "axis") == 0 self.sympy_data_[node.output[0]] = [] @@ -895,7 +911,7 @@ class SymbolicShapeInference: self.sympy_data_[node.output[0]] = numpy_helper.to_array(t) def _infer_ConstantOfShape(self, node): # noqa: N802 - sympy_shape = self._get_int_values(node)[0] + sympy_shape = self._get_int_or_float_values(node)[0] vi = self.known_vi_[node.output[0]] if sympy_shape is not None: if type(sympy_shape) != list: @@ -1442,7 +1458,7 @@ class SymbolicShapeInference: def _infer_Range(self, node): # noqa: N802 vi = self.known_vi_[node.output[0]] - input_data = self._get_int_values(node) + input_data = self._get_int_or_float_values(node) if all([i is not None for i in input_data]): start = as_scalar(input_data[0]) limit = as_scalar(input_data[1]) @@ -1496,7 +1512,7 @@ class SymbolicShapeInference: axes = get_attribute(node, "axes") keep_dims = get_attribute(node, "keepdims", 1) if keep_dims == 0 and axes == [0]: - data = self._get_int_values(node)[0] + data = self._get_int_or_float_values(node)[0] if data is not None: self.sympy_data_[node.output[0]] = sympy_reduce_product(data) @@ -1964,7 +1980,7 @@ class SymbolicShapeInference: if get_opset(self.out_mp_) <= 9: k = get_attribute(node, "k") else: - k = self._get_int_values(node)[1] + k = self._get_int_or_float_values(node)[1] if k is None: k = self._new_symbolic_dim_from_output(node) diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 18fe6517da..d0ea59b994 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -3,9 +3,11 @@ import os +import numpy + # -*- coding: UTF-8 -*- import onnx -from onnx import AttributeProto, GraphProto, TensorProto, helper # noqa: F401 +from onnx import AttributeProto, GraphProto, TensorProto, helper, numpy_helper # noqa: F401 if os.path.exists( os.path.join( @@ -340,6 +342,56 @@ class TestSymbolicShapeInferenceForOperators(unittest.TestCase): def test_einsum_transpose(self): self._test_einsum_one_input_impl(["a", "b"], ["b", "a"], "ij -> ji") + def test_mul_precision(self): + graph_input = onnx.helper.make_tensor_value_info("input", TensorProto.FLOAT, [1024]) + graph_output = onnx.helper.make_tensor_value_info("output", TensorProto.FLOAT, None) + + # initializers + value = numpy.array([0.5], dtype=numpy.float32) + constant = numpy_helper.from_array(value, name="constant") + + nodes = [ + # Get the shape of the input tensor: `input_tensor_shape = [1024]`. + onnx.helper.make_node("Shape", ["input"], ["input_shape"]), + # mul(1024, 0.5) => 512 + onnx.helper.make_node("Mul", ["input_shape", "constant"], ["output_shape"]), + # Resize input + onnx.helper.make_node( + "Resize", inputs=["input", "", "", "output_shape"], outputs=["output"], mode="nearest" + ), + ] + + graph_def = onnx.helper.make_graph(nodes, "TestMulPrecision", [graph_input], [graph_output], [constant]) + model = SymbolicShapeInference.infer_shapes(onnx.helper.make_model(graph_def)) + output_dims = unique_element(model.graph.output).type.tensor_type.shape.dim + self.assertEqual(len(output_dims), 1) + self.assertEqual(output_dims[0].dim_value, 512) + + def test_div_precision(self): + graph_input = onnx.helper.make_tensor_value_info("input", TensorProto.FLOAT, [768]) + graph_output = onnx.helper.make_tensor_value_info("output", TensorProto.FLOAT, None) + + # initializers + value = numpy.array([1.5], dtype=numpy.float32) + constant = numpy_helper.from_array(value, name="constant") + + nodes = [ + # Get the shape of the input tensor: `input_tensor_shape = [768]`. + onnx.helper.make_node("Shape", ["input"], ["input_shape"]), + # div(768, 1.5) => 512 + onnx.helper.make_node("Div", ["input_shape", "constant"], ["output_shape"]), + # Resize input + onnx.helper.make_node( + "Resize", inputs=["input", "", "", "output_shape"], outputs=["output"], mode="nearest" + ), + ] + + graph_def = onnx.helper.make_graph(nodes, "TestDivPrecision", [graph_input], [graph_output], [constant]) + model = SymbolicShapeInference.infer_shapes(onnx.helper.make_model(graph_def)) + output_dims = unique_element(model.graph.output).type.tensor_type.shape.dim + self.assertEqual(len(output_dims), 1) + self.assertEqual(output_dims[0].dim_value, 512) + class TestSymbolicShapeInferenceForSlice(unittest.TestCase): def check_slice_of_concat(self, input_dims, start, end, step, expected_output_dim):