diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index a3e6140381a0..be4dd7da903d 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3821,6 +3821,13 @@ def resize3d_dyn(d, r, s0, s1, s2): class Resize(OnnxOpConverter): """Converts an onnx Resize node into an equivalent Relax expression.""" + @classmethod + def _impl_v10(cls, bb, inputs, attr, params): + # Resize-10 takes (X, scales) and has the same semantics as Upsample-9. + scales = _get_constant_scales(inputs[1], params, "Resize") + mode = attr.get("mode", b"nearest") + return _legacy_upsample(bb, inputs[0], scales, mode, "Resize") + @classmethod def _impl_v18(cls, bb, inputs, attr, params): # Extract the many attributes of resize. @@ -3950,6 +3957,9 @@ def _impl_v18(cls, bb, inputs, attr, params): extrapolation_value=extrapolation_value, ) + # Resize-11 through Resize-18 share the (X, roi, scales, sizes) signature. + _impl_v11 = _impl_v18 + class AffineGrid(OnnxOpConverter): """Converts an onnx AffineGrid node into an equivalent Relax expression.""" @@ -5456,33 +5466,98 @@ def _impl_v9(cls, bb, inputs, attr, params): ) -class Upsample(OnnxOpConverter): - """Operator converter for Upsample (nearest mode).""" +def _legacy_upsample(bb, data, scales, mode, op_name): + """Lower Upsample (opset 7-9) and Resize-10, which share the same semantics. - @classmethod - def _impl_v9(cls, bb, inputs, attr, params): - scales = attr.get("scales") - assert len(scales) == 4 - assert scales[0] == scales[1] == 1 + Each output extent is ``floor(input * scale)`` and coordinates are mapped with + ``x_in = x_out / scale`` (asymmetric). In nearest mode the source index is + rounded down when upsampling and up when downsampling. + """ + mode = mode.decode("ascii") if isinstance(mode, bytes) else mode + if mode not in ("nearest", "linear"): + raise tvm.error.OpAttributeInvalid( + f'Value {mode} in attribute "mode" of operator {op_name} is not valid.' + ) + scales = [float(s) for s in scales] + shape = list(data.ty.shape) + ndims = len(shape) + if len(scales) != ndims: + raise tvm.error.OpAttributeInvalid( + f"{op_name} expects {ndims} scales for a rank-{ndims} input, got {len(scales)}." + ) + if ndims not in (3, 4, 5) or scales[0] != 1.0 or scales[1] != 1.0: + raise tvm.error.OpAttributeUnImplemented( + f"{op_name} is only supported for 3-D/4-D/5-D inputs with unit batch and " + f"channel scales, got scales {scales} for a rank-{ndims} input." + ) - inp_shape = [int(x) for x in inputs[0].ty.shape] - assert len(inp_shape) == 4 - out_shape2d = [int(dim * scale) for dim, scale in zip(inp_shape[2:], scales[2:])] + sizes = [] + for dim, scale in zip(shape[2:], scales[2:]): + if isinstance(dim, tirx.IntImm): + sizes.append(math.floor(int(dim) * scale)) + elif scale.is_integer(): + sizes.append(dim * int(scale)) + else: + sizes.append((dim.astype("float32") * scale).astype("int64")) + + spatial_scales = scales[2:] + if mode == "linear": + method, rounding_method = "linear", "" + elif all(s >= 1.0 for s in spatial_scales): + method, rounding_method = "nearest_neighbor", "floor" + elif all(s < 1.0 for s in spatial_scales): + method, rounding_method = "nearest_neighbor", "ceil" + else: + raise tvm.error.OpAttributeUnImplemented( + f"{op_name} nearest mode cannot mix upsampling and downsampling axes, " + f"got scales {scales}." + ) - mode = attr.get("mode", b"nearest").decode("ascii") - if mode == "nearest": - mode = "nearest_neighbor" - msg = f'Value {mode} in attribute "mode" of operator Upsample is not valid.' - assert mode in ("linear", "nearest_neighbor", "cubic"), msg + if ndims == 3: + return bb.emit_te( + topi.image.resize1d, + data, + [0.0, 0.0], + sizes, + "NCW", + method, + "asymmetric", + rounding_method, + ) + resize_op = relax.op.image.resize2d if ndims == 4 else relax.op.image.resize3d + return resize_op( + data, + size=relax.ShapeExpr(sizes), + layout="NCHW" if ndims == 4 else "NCDHW", + method=method, + coordinate_transformation_mode="asymmetric", + rounding_method=rounding_method, + ) - return relax.op.image.resize2d( - data=inputs[0], - roi=None, - size=relax.ShapeExpr(out_shape2d), # (H, W) - layout="NCHW", - method=mode, - coordinate_transformation_mode="asymmetric", # Align with Upsample + +def _get_constant_scales(scales, params, op_name): + scales = get_constant(scales, params) + if not isinstance(scales, tvm.ir.GenericConst): + raise tvm.error.OpAttributeUnImplemented( + f"{op_name} with non-constant scales is not supported." ) + return scales.value.numpy().tolist() + + +class Upsample(OnnxOpConverter): + """Converts an onnx Upsample node into an equivalent Relax expression.""" + + @classmethod + def _impl_v7(cls, bb, inputs, attr, params): + mode = attr.get("mode", b"nearest") + return _legacy_upsample(bb, inputs[0], attr["scales"], mode, "Upsample") + + @classmethod + def _impl_v9(cls, bb, inputs, attr, params): + # Since opset 9 the scales are an input rather than an attribute. + scales = _get_constant_scales(inputs[1], params, "Upsample") + mode = attr.get("mode", b"nearest") + return _legacy_upsample(bb, inputs[0], scales, mode, "Upsample") class HardSigmoid(OnnxOpConverter): diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index d7aa987c0527..6caaa2b849a6 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -9432,6 +9432,97 @@ def _visit(expr): assert seen_resize3d +def _make_legacy_upsample_model(op_name, opset, input_shape, scales, mode): + """Upsample-7 takes scales as an attribute; Upsample-9 and Resize-10 take an input.""" + inputs, initializer, attrs = ["X"], [], {"mode": mode} + if op_name == "Upsample" and opset < 9: + attrs["scales"] = [float(s) for s in scales] + else: + inputs.append("scales") + initializer.append(helper.make_tensor("scales", TensorProto.FLOAT, [len(scales)], scales)) + node = helper.make_node(op_name, inputs, ["Y"], **attrs) + graph = helper.make_graph( + [node], + "legacy_upsample", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], + initializer=initializer, + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)], + ) + return helper.make_model(graph, producer_name="legacy_upsample") + + +@pytest.mark.parametrize("op_name, opset", [("Upsample", 7), ("Upsample", 9), ("Resize", 10)]) +@pytest.mark.parametrize( + "input_shape, scales, mode", + [ + ([1, 2, 4, 5], [1.0, 1.0, 2.0, 2.0], "nearest"), + ([1, 2, 4, 5], [1.0, 1.0, 3.0, 2.0], "nearest"), + # floor(5 * 2.5) = 12: the nearest source index uses the given scale. + ([1, 2, 4, 5], [1.0, 1.0, 1.5, 2.5], "nearest"), + ([1, 2, 4, 5], [1.0, 1.0, 2.0, 3.0], "linear"), + ([1, 2, 5], [1.0, 1.0, 2.0], "nearest"), + ([1, 2, 3, 4, 5], [1.0, 1.0, 2.0, 1.0, 3.0], "nearest"), + ([1, 2, 3, 4, 5], [1.0, 1.0, 2.0, 2.0, 2.0], "linear"), + ], +) +def test_legacy_upsample(op_name, opset, input_shape, scales, mode): + model = _make_legacy_upsample_model(op_name, opset, input_shape, scales, mode) + check_correctness(model, opset=opset) + + +@pytest.mark.parametrize("mode", ["nearest", "linear"]) +def test_resize_v10_downsample(mode): + # Resize-10 allows scales below one; nearest then rounds the source index up. + model = _make_legacy_upsample_model("Resize", 10, [1, 2, 6, 5], [1.0, 1.0, 0.5, 0.6], mode) + check_correctness(model, opset=10) + + +@pytest.mark.parametrize("op_name, opset", [("Upsample", 7), ("Upsample", 9), ("Resize", 10)]) +def test_legacy_upsample_symbolic_shape(op_name, opset): + model = _make_legacy_upsample_model( + op_name, opset, ["N", 2, "H", "W"], [1.0, 1.0, 2.0, 1.5], "nearest" + ) + func = from_onnx(model, opset=opset, keep_params_in_input=True)["main"] + n, _, h, _ = func.params[0].ty.shape.values + out_n, out_c, out_h, _ = func.ret_ty.shape.values + tvm.ir.assert_structural_equal(out_n, n) + assert int(out_c) == 2 + # Integer scales keep the output extent an exact integer expression. + assert tvm.sym.Analyzer().can_prove_equal(out_h, h * 2) + + x = generate_random_value([2, 2, 4, 6], TensorProto.FLOAT) + check_correctness(model, inputs={"X": x}, opset=opset) + + +def test_resize_v11_still_uses_roi_scales_sizes_signature(): + # Adding Resize-10 must not change which converter handles opsets 11-17. + node = helper.make_node("Resize", ["X", "roi", "scales"], ["Y"], mode="nearest") + graph = helper.make_graph( + [node], + "resize_v11", + inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 5])], + initializer=[ + helper.make_tensor("roi", TensorProto.FLOAT, [0], []), + helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0]), + ], + outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)], + ) + model = helper.make_model(graph, producer_name="resize_v11") + check_correctness(model, opset=11) + + +def test_legacy_upsample_unsupported_scales(): + model = _make_legacy_upsample_model( + "Upsample", 9, [1, 2, 4, 5], [1.0, 2.0, 2.0, 2.0], "nearest" + ) + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="unit batch and channel"): + from_onnx(model, opset=9) + + model = _make_legacy_upsample_model("Resize", 10, [1, 2, 4, 5], [1.0, 1.0, 0.5, 2.0], "nearest") + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="cannot mix"): + from_onnx(model, opset=10) + + def test_einsum(): eqn = "ij->i" einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)