This is an automated email from the ASF dual-hosted git repository. guan404ming pushed a commit to branch fix/onnx-affinegrid-align-corners in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 5cb3fe1bb113f00d74f5c44418e7acab3a6145ce Author: Guan-Ming (Wesley) Chiu <[email protected]> AuthorDate: Fri Jun 19 23:36:06 2026 +0800 [Relax][ONNX] Support align_corners in AffineGrid op --- include/tvm/relax/attrs/image.h | 13 ++++++++++++ python/tvm/relax/frontend/onnx/onnx_frontend.py | 7 ++----- python/tvm/relax/op/image/image.py | 13 ++++++------ python/tvm/relax/transform/legalize_ops/image.py | 1 + python/tvm/topi/image/grid_sample.py | 25 ++++++++++++++++++------ src/relax/op/image/resize.cc | 9 +++++++-- src/relax/op/image/resize.h | 2 +- tests/python/relax/test_frontend_onnx.py | 12 +++++------- 8 files changed, 54 insertions(+), 28 deletions(-) diff --git a/include/tvm/relax/attrs/image.h b/include/tvm/relax/attrs/image.h index eacbea7180..c9a7203740 100644 --- a/include/tvm/relax/attrs/image.h +++ b/include/tvm/relax/attrs/image.h @@ -149,6 +149,19 @@ struct GridSampleAttrs : public AttrsNode { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode); }; // struct GridSampleAttrs +/*! \brief Attributes used in image affine_grid operator */ +struct AffineGridAttrs : public AttrsNode { + bool align_corners; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef<AffineGridAttrs>().def_ro( + "align_corners", &AffineGridAttrs::align_corners, + "If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers."); + } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode); +}; // struct AffineGridAttrs + } // namespace relax } // namespace tvm diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index d550d3bc00..61f95e7130 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3309,10 +3309,7 @@ class AffineGrid(OnnxOpConverter): def _impl_v20(cls, bb, inputs, attr, params): theta = inputs[0] # [N, 2, 3] for 2D size = get_constant(inputs[1], params) # [N, C, H, W] for 2D - align_corners = attr.get("align_corners", 0) - - if align_corners != 1: - raise NotImplementedError("AffineGrid with align_corners=0 is not yet supported in TVM") + align_corners = bool(attr.get("align_corners", 0)) # Extract size values if isinstance(size, relax.Constant): @@ -3328,7 +3325,7 @@ class AffineGrid(OnnxOpConverter): target_h, target_w = size_vals[2], size_vals[3] # Relax affine_grid outputs [N, 2, H, W] - grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w))) + grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w), align_corners)) # Permute to ONNX convention [N, H, W, 2] return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1])) diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py index 323bfa74b5..29aa8457d2 100644 --- a/python/tvm/relax/op/image/image.py +++ b/python/tvm/relax/op/image/image.py @@ -237,6 +237,7 @@ def grid_sample( def affine_grid( data: Expr, size: Expr | SizeLike, + align_corners: bool = True, ) -> Expr: """Generate a 2D sampling grid using an affine transformation matrix. @@ -253,20 +254,18 @@ def affine_grid( The target output spatial shape (H, W). If a single integer or PrimExpr is provided, it is interpreted as a square output shape (size, size). + align_corners : bool + If True, normalized grid coordinates map to corner pixels; if False, to + pixel centers (the PyTorch / ONNX default). + Returns ------- result : relax.Expr The output grid tensor with shape [batch, 2, H, W]. - - Note - ---- - Only `align_corners=True` is supported by this operator, matching the - behavior of the underlying TOPI implementation. When using this operator - via PyTorch or ONNX frontends, `align_corners=False` will be rejected. """ if isinstance(size, int | PrimExpr): size = (size, size) if isinstance(size, tuple | list): size = ShapeExpr(size) - return cast(Expr, _ffi_api.affine_grid(data, size)) + return cast(Expr, _ffi_api.affine_grid(data, size, align_corners)) diff --git a/python/tvm/relax/transform/legalize_ops/image.py b/python/tvm/relax/transform/legalize_ops/image.py index 42cb72f7c9..009c9a4512 100644 --- a/python/tvm/relax/transform/legalize_ops/image.py +++ b/python/tvm/relax/transform/legalize_ops/image.py @@ -66,6 +66,7 @@ def _image_affine_grid(bb: BlockBuilder, call: Call) -> Expr: topi.image.affine_grid, call.args[0], target_shape=target_shape, + align_corners=bool(call.attrs.align_corners), ) diff --git a/python/tvm/topi/image/grid_sample.py b/python/tvm/topi/image/grid_sample.py index 79032f41b3..d43d4a98fa 100644 --- a/python/tvm/topi/image/grid_sample.py +++ b/python/tvm/topi/image/grid_sample.py @@ -20,7 +20,7 @@ from tvm import te, tirx -def affine_grid(data, target_shape): +def affine_grid(data, target_shape, align_corners=True): """affine_grid operator that generates 2D sampling grid. This operation is described in https://arxiv.org/pdf/1506.02025.pdf. It generates a uniform @@ -35,6 +35,10 @@ def affine_grid(data, target_shape): target_shape: list/tuple of two int Specifies the output shape (H, W). + align_corners : bool + If True, normalized coordinates map to corner pixels; if False, to pixel centers + (the PyTorch / ONNX default). + Returns ------- Output : tvm.Tensor @@ -47,13 +51,22 @@ def affine_grid(data, target_shape): ) dtype = data.dtype - y_step = tirx.const((2.0 - 1e-7) / (target_shape[0] - 1), dtype=dtype) - x_step = tirx.const((2.0 - 1e-7) / (target_shape[1] - 1), dtype=dtype) - start = tirx.const(-1.0, dtype=dtype) + height, width = target_shape[0], target_shape[1] + if align_corners: + y_step = tirx.const((2.0 - 1e-7) / (height - 1), dtype=dtype) + x_step = tirx.const((2.0 - 1e-7) / (width - 1), dtype=dtype) + y_start = tirx.const(-1.0, dtype=dtype) + x_start = tirx.const(-1.0, dtype=dtype) + else: + # Pixel centers: coordinate i maps to (2 * i + 1) / size - 1. + y_step = tirx.const(2.0 / height, dtype=dtype) + x_step = tirx.const(2.0 / width, dtype=dtype) + y_start = tirx.const(-1.0 + 1.0 / height, dtype=dtype) + x_start = tirx.const(-1.0 + 1.0 / width, dtype=dtype) def _compute(n, dim, i, j): - y = start + i * y_step - x = start + j * x_step + y = y_start + i * y_step + x = x_start + j * x_step return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2] oshape = (data.shape[0], len(target_shape), *target_shape) diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index beb89af087..b92167e031 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -354,11 +354,15 @@ TVM_REGISTER_OP("relax.image.grid_sample") /* relax.image.affine_grid */ -Expr affine_grid(Expr data, Expr size) { +Expr affine_grid(Expr data, Expr size, bool align_corners) { + ffi::ObjectPtr<AffineGridAttrs> attrs = ffi::make_object<AffineGridAttrs>(); + attrs->align_corners = align_corners; static const Op& op = Op::Get("relax.image.affine_grid"); - return Call(op, {std::move(data), std::move(size)}, Attrs(), {}); + return Call(op, {std::move(data), std::move(size)}, Attrs(attrs), {}); } +TVM_FFI_STATIC_INIT_BLOCK() { AffineGridAttrs::RegisterReflection(); } + TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("relax.op.image.affine_grid", affine_grid); @@ -438,6 +442,7 @@ TVM_REGISTER_OP("relax.image.affine_grid") .set_num_inputs(2) .add_argument("data", "Tensor", "The input affine matrix tensor.") .add_argument("size", "Shape", "The target output shape (H, W).") + .set_attrs_type<AffineGridAttrs>() .set_attr<FInferType>("FInferType", InferTypeAffineGrid) .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow) .set_attr<bool>("FPurity", true); diff --git a/src/relax/op/image/resize.h b/src/relax/op/image/resize.h index 06a927d3a7..382a3a162b 100644 --- a/src/relax/op/image/resize.h +++ b/src/relax/op/image/resize.h @@ -49,7 +49,7 @@ Expr grid_sample(Expr data, Expr grid, ffi::String method, ffi::String layout, ffi::String padding_mode, bool align_corners); /*! \brief Image affine_grid operator. */ -Expr affine_grid(Expr data, Expr size); +Expr affine_grid(Expr data, Expr size, bool align_corners); } // namespace relax } // namespace tvm diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index d3036a2547..82e3f997d2 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -5555,13 +5555,11 @@ def test_nms_score_threshold(): ) -def test_affine_grid(): - affine_grid_node = helper.make_node( - "AffineGrid", - ["theta", "size"], - ["grid"], - align_corners=1, - ) +# align_corners=None omits the attribute, exercising the ONNX default of 0. [email protected]("align_corners", [None, 0, 1]) +def test_affine_grid(align_corners): + attrs = {} if align_corners is None else {"align_corners": align_corners} + affine_grid_node = helper.make_node("AffineGrid", ["theta", "size"], ["grid"], **attrs) graph = helper.make_graph( [affine_grid_node],
