diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 0e3ccef08cf0..ce92975cdeeb 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3113,6 +3113,7 @@ def _emit_resize_topi_dynamic_roi( cubic_coeff_a: float, exclude_outside: int, extrapolation_value: float, + scales: list | None, ) -> relax.Expr: """Lower Resize with runtime ROI via TOPI, which supports Expr ROI.""" if rank == 3: @@ -3129,6 +3130,7 @@ def resize1d_dyn(d, r, s0): cubic_coeff_a, exclude_outside, extrapolation_value, + scales=scales, ) return bb.emit_te(resize1d_dyn, data, roi_spatial_vec, sizes_spatial[0]) @@ -3147,12 +3149,14 @@ def resize2d_dyn(d, r, s0, s1): bicubic_alpha=cubic_coeff_a, bicubic_exclude=exclude_outside, extrapolation_value=extrapolation_value, + scales=scales, ) return bb.emit_te(resize2d_dyn, data, roi_spatial_vec, sizes_spatial[0], sizes_spatial[1]) def resize3d_dyn(d, r, s0, s1, s2): # r is ONNX order (D,H,W) x2; TOPI expects (W,H,D) x2. + # NOTE: scales order must stay ONNX (D,H,W), unlike r which is reordered return topi.image.resize3d( d, (r[2], r[1], r[0], r[5], r[4], r[3]), @@ -3164,6 +3168,7 @@ def resize3d_dyn(d, r, s0, s1, s2): bicubic_alpha=cubic_coeff_a, bicubic_exclude=exclude_outside, extrapolation_value=extrapolation_value, + scales=scales, ) return bb.emit_te( @@ -3230,7 +3235,10 @@ def _impl_v18(cls, bb, inputs, attr, params): use_dynamic_roi = roi_dynamic_vec is not None - # Convert scales to sizes if needed. + # Convert scales to sizes if needed, preserving the original spatial scales so + # the coordinate transformation uses the exact ONNX scale value rather than the + # lossy ratio derived from floor(input * scale) / input. + original_spatial_scales = None if scales is not None: if isinstance(scales, relax.Constant): scales = scales.data.numpy() @@ -3238,6 +3246,7 @@ def _impl_v18(cls, bb, inputs, attr, params): scales = [int(val.value) for val in scales.values] else: raise ValueError(f"Type {type(scales)} for scale is currently unsupported.") + original_spatial_scales = [float(s) for s in scales[2:]] sizes = [] for i, dim in enumerate(x.struct_info.shape): @@ -3264,6 +3273,7 @@ def _impl_v18(cls, bb, inputs, attr, params): cubic_coeff_a, exclude_outside, extrapolation_value, + original_spatial_scales, ) if ndims == 3: @@ -3279,8 +3289,24 @@ def _impl_v18(cls, bb, inputs, attr, params): cubic_coeff_a, exclude_outside, extrapolation_value, + scales=original_spatial_scales, ) elif ndims == 4: + if original_spatial_scales is not None: + return bb.emit_te( + topi.image.resize2d, + x, + roi_static, + sizes, + "NCHW", + topi_mode, + coord_mode, + rounding_method, + cubic_coeff_a, + exclude_outside, + extrapolation_value, + scales=original_spatial_scales, + ) return relax.op.image.resize2d( x, size=relax.ShapeExpr(sizes), @@ -3295,6 +3321,21 @@ def _impl_v18(cls, bb, inputs, attr, params): ) else: # ndims == 5 roi3d = _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_static) + if original_spatial_scales is not None: + return bb.emit_te( + topi.image.resize3d, + x, + roi3d, + sizes, + "NCDHW", + topi_mode, + coord_mode, + rounding_method, + cubic_coeff_a, + exclude_outside, + extrapolation_value, + scales=original_spatial_scales, + ) return relax.op.image.resize3d( x, size=relax.ShapeExpr(sizes), diff --git a/python/tvm/topi/image/resize.py b/python/tvm/topi/image/resize.py index 1f4799c8ecc8..58c9d1baffea 100644 --- a/python/tvm/topi/image/resize.py +++ b/python/tvm/topi/image/resize.py @@ -145,9 +145,13 @@ def get_inx( start_x=0, end_x=-1, use_int_div=False, + scale_x_override=None, ): """Infer input x from output x with various coordinate transformation methods""" - scale_x = te.div(image_width.astype("float"), target_width.astype("float")) + if scale_x_override is not None: + scale_x = scale_x_override + else: + scale_x = te.div(image_width.astype("float"), target_width.astype("float")) if coordinate_transformation_mode == "half_pixel": in_x = (x + 0.5) * scale_x - 0.5 elif coordinate_transformation_mode == "align_corners": @@ -237,6 +241,7 @@ def _resize_1d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_x=None, ): """Perform resize operation on the data with selected method and options. @@ -315,7 +320,15 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") - in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[1]) + in_x = get_inx( + x, + image_width, + target_width, + coordinate_transformation_mode, + roi[0], + roi[1], + scale_x_override=scale_x, + ) if method == "nearest_neighbor": if rounding_method == "": @@ -383,6 +396,7 @@ def resize1d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -472,6 +486,8 @@ def resize1d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_x = (1.0 / scales[0]) if scales is not None else None + def compute_func(*indices): return _resize_1d( indices, @@ -487,6 +503,7 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_x=scale_x, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) @@ -510,6 +527,8 @@ def _resize_2d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_h=None, + scale_w=None, ): """Perform resize operation on the data with selected method and options. @@ -593,8 +612,10 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): width_use_int_div = False if method == "nearest_neighbor" and coordinate_transformation_mode == "asymmetric": if rounding_method == "floor" or rounding_method == "": - height_use_int_div = can_convert_multiply_to_intdiv(image_height, target_height) - width_use_int_div = can_convert_multiply_to_intdiv(image_width, target_width) + if scale_h is None: + height_use_int_div = can_convert_multiply_to_intdiv(image_height, target_height) + if scale_w is None: + width_use_int_div = can_convert_multiply_to_intdiv(image_width, target_width) n, c, y, x, cc, inum, ic = get_2d_indices(indices, layout) box_idx = box_indices(n) if box_indices is not None else n @@ -618,6 +639,7 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): roi[1], roi[3], width_use_int_div, + scale_x_override=scale_w, ) in_y = get_inx( y, @@ -627,6 +649,7 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): roi[0], roi[2], height_use_int_div, + scale_x_override=scale_h, ) if method == "nearest_neighbor": @@ -756,6 +779,7 @@ def resize2d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -839,6 +863,9 @@ def resize2d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_h = (1.0 / scales[0]) if scales is not None else None + scale_w = (1.0 / scales[1]) if scales is not None else None + def compute_func(*indices): return _resize_2d( indices, @@ -856,6 +883,8 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_h=scale_h, + scale_w=scale_w, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) @@ -976,6 +1005,9 @@ def _resize_3d( alpha=-0.5, exclude_outside=0, out_dtype=None, + scale_d=None, + scale_h=None, + scale_w=None, ): """Perform resize operation on the data with selected method and options. @@ -1066,9 +1098,33 @@ def _cast_output(value, data_dtype="float32", out_dtype=None): if boxes is not None: # TODO(mbrookhart): Find an example of this raise NotImplementedError("resize1d with image boxes not yet implemented") - in_z = get_inx(z, image_depth, target_depth, coordinate_transformation_mode, roi[2], roi[5]) - in_y = get_inx(y, image_height, target_height, coordinate_transformation_mode, roi[1], roi[4]) - in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode, roi[0], roi[3]) + in_z = get_inx( + z, + image_depth, + target_depth, + coordinate_transformation_mode, + roi[2], + roi[5], + scale_x_override=scale_d, + ) + in_y = get_inx( + y, + image_height, + target_height, + coordinate_transformation_mode, + roi[1], + roi[4], + scale_x_override=scale_h, + ) + in_x = get_inx( + x, + image_width, + target_width, + coordinate_transformation_mode, + roi[0], + roi[3], + scale_x_override=scale_w, + ) if method == "nearest_neighbor": if rounding_method == "": @@ -1225,6 +1281,7 @@ def resize3d( extrapolation_value=0.0, out_dtype=None, output_shape=None, + scales=None, ): """Perform resize operation on the data. @@ -1302,6 +1359,10 @@ def resize3d( if isinstance(size[i], int): size[i] = tvm.tirx.IntImm("int32", size[i]) + scale_d = (1.0 / scales[0]) if scales is not None else None + scale_h = (1.0 / scales[1]) if scales is not None else None + scale_w = (1.0 / scales[2]) if scales is not None else None + def compute_func(*indices): return _resize_3d( indices, @@ -1321,6 +1382,9 @@ def compute_func(*indices): exclude_outside=bicubic_exclude, extrapolation_value=extrapolation_value, out_dtype=out_dtype, + scale_d=scale_d, + scale_h=scale_h, + scale_w=scale_w, ) return te.compute(output_shape, compute_func, name="resize", tag=tag.INJECTIVE) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 471186589e75..2774bb84e1af 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -4012,6 +4012,203 @@ def _visit(expr): assert seen_resize3d +@pytest.mark.parametrize( + "coord_mode, method", + [ + ("half_pixel", "nearest"), + ("pytorch_half_pixel", "nearest"), + ("asymmetric", "nearest"), + ("half_pixel", "linear"), + ], +) +def test_resize_noninteger_scales_2d(coord_mode, method): + """Non-integer scales must use the original scale in coordinate transformation. + + floor(3 * 2.5) = 7, so the recomputed ratio 3/7 = 0.4286 differs from 1/2.5 = 0.4, + causing wrong pixel mapping at boundary positions before the fix. + """ + nearest_mode_kwargs = {} + if method == "nearest": + nearest_mode_kwargs["nearest_mode"] = "round_prefer_floor" + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode=method, + coordinate_transformation_mode=coord_mode, + **nearest_mode_kwargs, + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.5, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_noninteger_scales_1d(): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="half_pixel", + nearest_mode="round_prefer_floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_1d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [3], [1.0, 1.0, 1.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_noninteger_scales_3d(): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_noninteger_3d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3, 3])], + initializer=[ + helper.make_tensor("scales", TensorProto.FLOAT, [5], [1.0, 1.0, 1.5, 1.5, 1.5]) + ], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_dynamic_roi_noninteger_scales_1d(): + resize_node = helper.make_node( + "Resize", + ["X", "roi", "scales"], + ["Y"], + mode="linear", + coordinate_transformation_mode="half_pixel", + ) + graph = helper.make_graph( + [resize_node], + "resize_dynamic_roi_noninteger_1d", + inputs=[ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3]), + helper.make_tensor_value_info("roi", TensorProto.FLOAT, [6]), + ], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [3], [1.0, 1.0, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_dynamic_roi_noninteger_scales_2d(): + resize_node = helper.make_node( + "Resize", + ["X", "roi", "scales"], + ["Y"], + mode="linear", + coordinate_transformation_mode="half_pixel", + ) + graph = helper.make_graph( + [resize_node], + "resize_dynamic_roi_noninteger_2d", + inputs=[ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 3]), + helper.make_tensor_value_info("roi", TensorProto.FLOAT, [8]), + ], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.5, 2.5])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 7, 7])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_asymmetric_nearest_noninteger_scales_2d(): + """Asymmetric+nearest+floor optimization must not ignore scale override. + + When coordinate_transformation_mode="asymmetric", method="nearest_neighbor", + rounding_method="floor", and scales is non-integer, the integer-division optimization + must not be applied, The bug: can_convert_multiply_to_intdiv checks only derived ratio + (ignoring scale override), causing wrong pixel mapping. + + Example: input 2x2, scale 2.4, output 4x4. Derived ratio 4/2=2.0 triggers optimization, + but floor(2*0.4167) != floor(2/2) at same coordinates. + """ + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_asymmetric_nearest_noninteger_scales_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.4, 2.4])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +def test_resize_asymmetric_nearest_integer_scales_baseline(): + """Baseline: asymmetric+nearest+floor with integer scale should work correctly. + + This ensures the guarded optimization (when scale_override is None) still produces + correct results for integer scales and doesn't regress existing behavior. + """ + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_asymmetric_nearest_integer_scales_2d", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2])], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0])], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4])], + ) + check_correctness(helper.make_model(graph), opset=18) + + +@pytest.mark.parametrize( + "input_shape,scales,output_shape", + [ + ([1, 1, 4, 4], [1.0, 1.0, 2.0, 2.0], [1, 1, 8, 8]), + ([1, 1, 3, 3], [1.0, 1.0, 3.0, 3.0], [1, 1, 9, 9]), + ], +) +def test_resize_integer_scales_regression(input_shape, scales, output_shape): + resize_node = helper.make_node( + "Resize", + ["X", "", "scales"], + ["Y"], + mode="nearest", + coordinate_transformation_mode="half_pixel", + nearest_mode="round_prefer_floor", + ) + graph = helper.make_graph( + [resize_node], + "resize_integer_scales", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], + initializer=[helper.make_tensor("scales", TensorProto.FLOAT, [len(scales)], scales)], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)], + ) + check_correctness(helper.make_model(graph), opset=18) + + def test_einsum(): eqn = "ij->i" einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)