Skip to content

Commit 2b4b58d

Browse files
authored
[Relax][Onnx][Resize] Handle non-4D input tensors (#18666)
### Motivation The ONNX Resize operator supports resizing N-D tensors per the ONNX specification. However, the current Relax ONNX frontend only supports 4D inputs and raises an assertion error for valid non-4D models. This PR extends the Relax ONNX Resize converter beyond the 4D-only restriction, aligning behavior with the ONNX specification and ONNX Runtime for supported ranks. ### Changes - Remove the 4D-only assertion in the Relax ONNX Resize converter - Preserve the existing 4D resize path without behavior changes - Support non-4D Resize by lowering to existing resize implementations for supported ranks - Ensure Resize attributes are handled correctly for non-4D cases ### Testing - Verified outputs against ONNX Runtime - Added ONNX Resize tests covering non-4D input cases (`test_resize_nd_sizes`) Fixes: [[Bug] Resize N-D import failure: TVM only supports 4D resize2d, but ONNX Resize supports N-D tensors https://github.com/apache/tvm/issues/18608](https://github.com/apache/tvm/issues/18608)
1 parent eb4eb3e commit 2b4b58d

2 files changed

Lines changed: 85 additions & 16 deletions

File tree

python/tvm/relax/frontend/onnx/onnx_frontend.py

Lines changed: 50 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2235,15 +2235,18 @@ def _impl_v18(cls, bb, inputs, attr, params):
22352235

22362236
# Adapt attributes to fit TVM definition.
22372237
if mode == "nearest":
2238-
mode = "nearest_neighbor"
2238+
relax_mode = "nearest_neighbor"
2239+
else:
2240+
relax_mode = mode
2241+
topi_mode = relax_mode
22392242

22402243
# Unpack inputs.
22412244
x = inputs[0]
22422245
roi = get_constant(inputs[1], params)
22432246
scales = get_constant(inputs[2], params)
22442247
sizes = get_constant(inputs[3], params)
22452248
ndims = len(x.struct_info.shape)
2246-
assert ndims == 4, "Only resize2d is currently supported."
2249+
assert ndims in (3, 4, 5), "Only resize1d/resize2d/resize3d are supported."
22472250

22482251
assert (
22492252
scales is None or sizes is None
@@ -2253,6 +2256,8 @@ def _impl_v18(cls, bb, inputs, attr, params):
22532256
if roi is not None:
22542257
if isinstance(roi, relax.Constant):
22552258
roi = roi.data.numpy().tolist()
2259+
if len(roi) == 2 * ndims:
2260+
roi = roi[2:ndims] + roi[ndims + 2 : 2 * ndims]
22562261
else:
22572262
roi = relax.op.concat(
22582263
[
@@ -2262,9 +2267,9 @@ def _impl_v18(cls, bb, inputs, attr, params):
22622267
axis=0,
22632268
)
22642269
# TODO The backend C++ func resize2d does not support dynamic ROI for now.
2265-
raise NotImplementedError("Dynamic ROI is not supported in resize2d for now.")
2270+
raise NotImplementedError("Dynamic ROI is not supported in resize for now.")
22662271
else:
2267-
roi = [0.0] * 4
2272+
roi = [0.0] * (2 * (ndims - 2))
22682273

22692274
# Convert scales to sizes if needed.
22702275
if scales is not None:
@@ -2287,18 +2292,47 @@ def _impl_v18(cls, bb, inputs, attr, params):
22872292
else:
22882293
assert f"Type {type(size)} for size is currently unsupported."
22892294

2290-
return relax.op.image.resize2d(
2291-
x,
2292-
size=relax.ShapeExpr(sizes),
2293-
roi=roi,
2294-
layout="NCHW",
2295-
method=mode,
2296-
coordinate_transformation_mode=coord_mode,
2297-
rounding_method=rounding_method,
2298-
cubic_alpha=cubic_coeff_a,
2299-
cubic_exclude=exclude_outside,
2300-
extrapolation_value=extrapolation_value,
2301-
)
2295+
if ndims == 3:
2296+
return bb.emit_te(
2297+
topi.image.resize1d,
2298+
x,
2299+
roi,
2300+
sizes,
2301+
"NCW",
2302+
topi_mode,
2303+
coord_mode,
2304+
rounding_method,
2305+
cubic_coeff_a,
2306+
exclude_outside,
2307+
extrapolation_value,
2308+
)
2309+
elif ndims == 4:
2310+
return relax.op.image.resize2d(
2311+
x,
2312+
size=relax.ShapeExpr(sizes),
2313+
roi=roi,
2314+
layout="NCHW",
2315+
method=relax_mode,
2316+
coordinate_transformation_mode=coord_mode,
2317+
rounding_method=rounding_method,
2318+
cubic_alpha=cubic_coeff_a,
2319+
cubic_exclude=exclude_outside,
2320+
extrapolation_value=extrapolation_value,
2321+
)
2322+
else: # ndims == 5
2323+
return bb.emit_te(
2324+
topi.image.resize3d,
2325+
x,
2326+
roi,
2327+
sizes,
2328+
"NCDHW",
2329+
topi_mode,
2330+
coord_mode,
2331+
rounding_method,
2332+
cubic_coeff_a,
2333+
exclude_outside,
2334+
extrapolation_value,
2335+
)
23022336

23032337

23042338
class Einsum(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2681,6 +2681,41 @@ def test_resize(with_roi, roi_list):
26812681
check_correctness(model)
26822682

26832683

2684+
def test_resize_nd_sizes():
2685+
cases = [
2686+
("resize1d", [1, 1, 4], [1, 1, 7]),
2687+
("resize2d", [1, 1, 4, 5], [1, 1, 6, 7]),
2688+
("resize3d", [1, 1, 3, 4, 5], [1, 1, 4, 6, 7]),
2689+
]
2690+
2691+
for name, input_shape, sizes in cases:
2692+
resize_node = helper.make_node(
2693+
"Resize",
2694+
["X", "", "", "sizes"],
2695+
["Y"],
2696+
mode="nearest",
2697+
coordinate_transformation_mode="asymmetric",
2698+
nearest_mode="floor",
2699+
)
2700+
2701+
graph = helper.make_graph(
2702+
[resize_node],
2703+
name,
2704+
inputs=[
2705+
helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape),
2706+
],
2707+
initializer=[
2708+
helper.make_tensor("sizes", TensorProto.INT64, [len(sizes)], sizes),
2709+
],
2710+
outputs=[
2711+
helper.make_tensor_value_info("Y", TensorProto.FLOAT, sizes),
2712+
],
2713+
)
2714+
2715+
model = helper.make_model(graph, producer_name=name)
2716+
check_correctness(model, opset=18)
2717+
2718+
26842719
def test_einsum():
26852720
eqn = "ij->i"
26862721
einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)

0 commit comments

Comments
 (0)