wuyii8941 opened a new issue, #19688:
URL: https://github.com/apache/tvm/issues/19688
## Expected behavior
ONNX `GridSample` (opset 16+) supports both 4D inputs (`[N, C, H, W]` with
grid `[N, H_out, W_out, 2]`) and 5D volumetric inputs (`[N, C, D, H, W]` with
grid `[N, D_out, H_out, W_out, 3]`). ONNX Runtime handles both.
## Actual behavior
The TVM ONNX frontend crashes immediately during conversion when the input
is 5D:
```
tvm.error.InternalError: PermuteDims expects the number of input axes to
equal
the ndim of the input tensor. However, the tensor ndim is 5 while the given
number of axes is 4
```
## Reproduction
```python
import numpy as np
import onnx
from onnx import helper, TensorProto
import onnxruntime as ort
from tvm.relax.frontend.onnx import from_onnx
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4, 4])
G = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 4, 4, 4, 3])
Y = onnx.ValueInfoProto(); Y.name = "Y"
node = helper.make_node("GridSample", ["X", "G"], ["Y"],
mode="nearest", padding_mode="zeros",
align_corners=0)
g = helper.make_graph([node], "g", [X, G], [Y])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 20)])
x = np.random.randn(1, 1, 4, 4, 4).astype(np.float32)
grid = np.random.uniform(-1.2, 1.2, size=(1, 4, 4, 4, 3)).astype(np.float32)
# ORT works
print("ORT shape:", ort.InferenceSession(m.SerializeToString()).run(None,
{"X": x, "G": grid})[0].shape)
# TVM crashes here
inf = onnx.shape_inference.infer_shapes(m)
mod = from_onnx(inf)
```
## Root cause
In `python/tvm/relax/frontend/onnx/onnx_frontend.py`, `GridSample._impl_v16`
unconditionally permutes the grid from `[N,H,W,2]` (ONNX) to `[N,2,H,W]` (TVM
NCHW) and calls the 2D `image.grid_sample`:
```python
class GridSample(OnnxOpConverter):
@classmethod
def _impl_v16(cls, bb, inputs, attr, params):
...
# ONNX grid shape: [N, H_out, W_out, 2]
# TVM grid shape: [N, 2, H_out, W_out]
grid = relax.op.permute_dims(grid, [0, 3, 1, 2]) # <-- hard-coded
4D
return relax.op.image.grid_sample(
data, grid, method=method, layout="NCHW", ... # <--
layout="NCHW" only
)
```
There is no dispatch on input rank, so 5D inputs fail at the `permute_dims`
step.
## Suggested fix
Dispatch on input rank and call the appropriate Relax op
(`image.grid_sample` handles 2D; volumetric grid sample would need to use the
3D variant or be added):
```python
ndim = data.struct_info.ndim
if ndim == 4:
grid = relax.op.permute_dims(grid, [0, 3, 1, 2])
return relax.op.image.grid_sample(data, grid, method=method,
layout="NCHW", ...)
elif ndim == 5:
grid = relax.op.permute_dims(grid, [0, 4, 1, 2, 3])
return relax.op.image.grid_sample(data, grid, method=method,
layout="NCDHW", ...)
else:
raise NotImplementedError(f"GridSample with ndim={ndim} is not
supported")
```
Note: the 5D path may require adding a `grid_sample_3d` Relax op + TOPI
lowering if it doesn't already exist; if not feasible immediately, the
converter should raise a clear `NotImplementedError` for the 5D case instead of
crashing inside `permute_dims`.
## Impact
Any 3D model using volumetric grid sample (3D medical imaging, video
stabilization, deformable 3D convs) cannot be imported.
## Environment
- TVM: latest `main` (commit b172d5ea3), reproduced against the same code in
upstream
- Python: 3.11
- ONNX Runtime: 1.24.4
--
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]
