guan404ming commented on code in PR #19863:
URL: https://github.com/apache/tvm/pull/19863#discussion_r3450375079
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -3322,15 +3322,17 @@ def _impl_v20(cls, bb, inputs, attr, params):
else:
raise NotImplementedError(f"Dynamic size of type {type(size)} is
not supported")
- # Only 2D is supported: size = [N, C, H, W]
- if len(size_vals) != 4:
- raise ValueError("Only 2D AffineGrid (size=[N,C,H,W]) is
supported")
- 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)))
- # Permute to ONNX convention [N, H, W, 2]
- return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
+ if len(size_vals) == 4:
+ # 2D: size = [N, C, H, W]; relax affine_grid outputs [N, 2, H, W].
+ grid = bb.emit(relax.op.image.affine_grid(theta,
tuple(size_vals[2:])))
+ # Permute to ONNX convention [N, H, W, 2].
+ return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
+ if len(size_vals) == 5:
+ # 3D: size = [N, C, D, H, W]; relax affine_grid outputs [N, 3, D,
H, W].
+ grid = bb.emit(relax.op.image.affine_grid(theta,
tuple(size_vals[2:])))
+ # Permute to ONNX convention [N, D, H, W, 3].
+ return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 4, 1]))
+ raise ValueError("AffineGrid expects size to be [N,C,H,W] (2D) or
[N,C,D,H,W] (3D)")
Review Comment:
Done, single emit with computed permute axes.
##########
python/tvm/topi/image/grid_sample.py:
##########
@@ -30,31 +30,38 @@ def affine_grid(data, target_shape):
Parameters
----------
data : tvm.Tensor
- 3-D with shape [batch, 2, 3]. The affine matrix.
+ 3-D with shape [batch, 2, 3] for 2D or [batch, 3, 4] for 3D. The
affine matrix.
- target_shape: list/tuple of two int
- Specifies the output shape (H, W).
+ target_shape: list/tuple of int
+ Specifies the output spatial shape (H, W) for 2D or (D, H, W) for 3D.
Returns
-------
Output : tvm.Tensor
- 4-D with shape [batch, 2, target_height, target_width]
+ [batch, 2, H, W] for 2D or [batch, 3, D, H, W] for 3D.
"""
assert target_shape is not None
- assert len(target_shape) == 2
- assert target_shape[0] > 1 and target_shape[1] > 1, (
- "target height/width should be greater than 1"
- )
+ assert len(target_shape) in (2, 3)
+ assert all(s > 1 for s in target_shape), "target spatial dims should be
greater than 1"
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)
+ steps = [tirx.const((2.0 - 1e-7) / (s - 1), dtype=dtype) for s in
target_shape]
+
+ if len(target_shape) == 2:
+
+ def _compute(n, dim, i, j):
+ y = start + i * steps[0]
+ x = start + j * steps[1]
+ return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
+
+ else:
- def _compute(n, dim, i, j):
- y = start + i * y_step
- x = start + j * x_step
- return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
+ def _compute(n, dim, k, i, j):
+ z = start + k * steps[0]
+ y = start + i * steps[1]
+ x = start + j * steps[2]
+ return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
* z + data[n, dim, 3]
Review Comment:
Done, replaced with a generic N-D compute.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
