Skip to content

Commit ec0daad

Browse files
authored
[Relax][ONNX] Support Resize dynamic ROI via TOPI (#18963)
The ONNX Resize converter previously rejected non-constant ROI inputs, which blocked models where ROI is provided at runtime. This change adds a dynamic-ROI path lowered through TOPI resize kernels while preserving the existing relax.image.resize* path for static ROI. Specifically: - add reusable helper to convert ONNX full ROI ([starts..., ends...]) into spatial ROI vector - add reusable helper to emit topi.image.resize1d/2d/3d for dynamic ROI - keep static ROI fast path for relax.image.resize2d/resize3d - normalize dynamic ROI expr before emit_te to ensure struct_info is populated - handle optional Resize inputs (roi/scales/sizes) more defensively - add frontend test coverage with graph-input ROI: test_resize_dynamic_roi_tf_crop_and_resize Ref: #18945
1 parent fd9d9db commit ec0daad

2 files changed

Lines changed: 199 additions & 23 deletions

File tree

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

Lines changed: 141 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2632,6 +2632,105 @@ def _impl_v1(cls, bb, inputs, attr, params):
26322632
return inputs[0]
26332633

26342634

2635+
def _onnx_resize_spatial_roi_vector(roi_full: relax.Expr, rank: int) -> relax.Expr:
2636+
"""Map ONNX ROI [starts..., ends...] to TOPI spatial ROI (drop N/C axes)."""
2637+
return relax.op.concat(
2638+
[
2639+
relax.op.strided_slice(roi_full, axes=[0], begin=[2], end=[rank]),
2640+
relax.op.strided_slice(roi_full, axes=[0], begin=[rank + 2], end=[2 * rank]),
2641+
],
2642+
axis=0,
2643+
)
2644+
2645+
2646+
def _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_spatial: list[float]) -> list[float]:
2647+
"""Reorder spatial ROI for NCDHW ONNX layout to TOPI resize3d convention.
2648+
2649+
ONNX spatial slice after dropping N/C is ordered (D, H, W) for starts then ends.
2650+
TOPI ``resize3d`` with layout NCDHW expects
2651+
``(start_w, start_h, start_d, end_w, end_h, end_d)`` (see topi/image/resize.py).
2652+
"""
2653+
if len(roi_spatial) != 6:
2654+
return roi_spatial
2655+
d0, h0, w0, d1, h1, w1 = roi_spatial
2656+
return [w0, h0, d0, w1, h1, d1]
2657+
2658+
2659+
def _emit_resize_topi_dynamic_roi(
2660+
bb: relax.BlockBuilder,
2661+
data: relax.Expr,
2662+
roi_spatial_vec: relax.Expr,
2663+
sizes_spatial: list,
2664+
rank: int,
2665+
topi_mode: str,
2666+
coord_mode: str,
2667+
rounding_method: str,
2668+
cubic_coeff_a: float,
2669+
exclude_outside: int,
2670+
extrapolation_value: float,
2671+
) -> relax.Expr:
2672+
"""Lower Resize with runtime ROI via TOPI, which supports Expr ROI."""
2673+
if rank == 3:
2674+
2675+
def resize1d_dyn(d, r, s0):
2676+
return topi.image.resize1d(
2677+
d,
2678+
(r[0], r[1]),
2679+
[s0],
2680+
"NCW",
2681+
topi_mode,
2682+
coord_mode,
2683+
rounding_method,
2684+
cubic_coeff_a,
2685+
exclude_outside,
2686+
extrapolation_value,
2687+
)
2688+
2689+
return bb.emit_te(resize1d_dyn, data, roi_spatial_vec, sizes_spatial[0])
2690+
2691+
if rank == 4:
2692+
2693+
def resize2d_dyn(d, r, s0, s1):
2694+
return topi.image.resize2d(
2695+
d,
2696+
(r[0], r[1], r[2], r[3]),
2697+
(s0, s1),
2698+
layout="NCHW",
2699+
method=topi_mode,
2700+
coordinate_transformation_mode=coord_mode,
2701+
rounding_method=rounding_method,
2702+
bicubic_alpha=cubic_coeff_a,
2703+
bicubic_exclude=exclude_outside,
2704+
extrapolation_value=extrapolation_value,
2705+
)
2706+
2707+
return bb.emit_te(resize2d_dyn, data, roi_spatial_vec, sizes_spatial[0], sizes_spatial[1])
2708+
2709+
def resize3d_dyn(d, r, s0, s1, s2):
2710+
# r is ONNX order (D,H,W) x2; TOPI expects (W,H,D) x2.
2711+
return topi.image.resize3d(
2712+
d,
2713+
(r[2], r[1], r[0], r[5], r[4], r[3]),
2714+
(s0, s1, s2),
2715+
layout="NCDHW",
2716+
method=topi_mode,
2717+
coordinate_transformation_mode=coord_mode,
2718+
rounding_method=rounding_method,
2719+
bicubic_alpha=cubic_coeff_a,
2720+
bicubic_exclude=exclude_outside,
2721+
extrapolation_value=extrapolation_value,
2722+
)
2723+
2724+
return bb.emit_te(
2725+
resize3d_dyn,
2726+
data,
2727+
roi_spatial_vec,
2728+
sizes_spatial[0],
2729+
sizes_spatial[1],
2730+
sizes_spatial[2],
2731+
)
2732+
2733+
26352734
class Resize(OnnxOpConverter):
26362735
"""Converts an onnx Resize node into an equivalent Relax expression."""
26372736

@@ -2654,36 +2753,39 @@ def _impl_v18(cls, bb, inputs, attr, params):
26542753

26552754
# Unpack inputs.
26562755
x = inputs[0]
2657-
roi = get_constant(inputs[1], params)
2658-
scales = get_constant(inputs[2], params)
2659-
sizes = get_constant(inputs[3], params)
2756+
roi = get_constant(inputs[1], params) if len(inputs) > 1 and inputs[1] is not None else None
2757+
scales = get_constant(inputs[2], params) if len(inputs) > 2 else None
2758+
sizes = get_constant(inputs[3], params) if len(inputs) > 3 else None
26602759
ndims = len(x.struct_info.shape)
26612760
assert ndims in (3, 4, 5), "Only resize1d/resize2d/resize3d are supported."
26622761

26632762
assert scales is None or sizes is None, (
26642763
"Only one of scales and sizes can be provided in Resize."
26652764
)
26662765

2667-
# Define relax implementation.
2766+
# ROI can be a static list (for relax.image.resize*) or dynamic tensor (TOPI path).
2767+
roi_static: list[float] | None = None
2768+
roi_dynamic_vec: relax.Expr | None = None
26682769
if roi is not None:
26692770
if isinstance(roi, relax.Constant):
2670-
roi = roi.data.numpy().tolist()
2671-
if len(roi) == 2 * ndims:
2672-
roi = roi[2:ndims] + roi[ndims + 2 : 2 * ndims]
2673-
elif len(roi) == 0:
2674-
roi = [0.0] * (2 * (ndims - 2))
2771+
roi_np = roi.data.numpy().tolist()
2772+
if len(roi_np) == 2 * ndims:
2773+
roi_static = roi_np[2:ndims] + roi_np[ndims + 2 : 2 * ndims]
2774+
elif len(roi_np) == 0:
2775+
roi_static = [0.0] * (2 * (ndims - 2))
2776+
elif len(roi_np) == 2 * (ndims - 2):
2777+
# Some exporters already provide spatial-only ROI.
2778+
roi_static = roi_np
2779+
else:
2780+
roi_static = roi_np
26752781
else:
2676-
roi = relax.op.concat(
2677-
[
2678-
relax.op.strided_slice(roi, axes=[0], begin=[2], end=[ndims]),
2679-
relax.op.strided_slice(roi, axes=[0], begin=[ndims + 2], end=[2 * ndims]),
2680-
],
2681-
axis=0,
2782+
roi_dynamic_vec = bb.normalize(
2783+
_onnx_resize_spatial_roi_vector(roi, ndims)
26822784
)
2683-
# TODO The backend C++ func resize2d does not support dynamic ROI for now.
2684-
raise NotImplementedError("Dynamic ROI is not supported in resize for now.")
26852785
else:
2686-
roi = [0.0] * (2 * (ndims - 2))
2786+
roi_static = [0.0] * (2 * (ndims - 2))
2787+
2788+
use_dynamic_roi = roi_dynamic_vec is not None
26872789

26882790
# Convert scales to sizes if needed.
26892791
if scales is not None:
@@ -2692,7 +2794,7 @@ def _impl_v18(cls, bb, inputs, attr, params):
26922794
elif isinstance(scales, relax.expr.ShapeExpr):
26932795
scales = [int(val.value) for val in scales.values]
26942796
else:
2695-
assert f"Type {type(scales)} for scale is currently unsupported."
2797+
raise ValueError(f"Type {type(scales)} for scale is currently unsupported.")
26962798
sizes = []
26972799

26982800
for i, dim in enumerate(x.struct_info.shape):
@@ -2704,13 +2806,28 @@ def _impl_v18(cls, bb, inputs, attr, params):
27042806
elif isinstance(sizes, relax.expr.ShapeExpr):
27052807
sizes = [int(val.value) for val in sizes.values][2:]
27062808
else:
2707-
assert f"Type {type(sizes)} for size is currently unsupported."
2809+
raise ValueError(f"Type {type(sizes)} for size is currently unsupported.")
2810+
2811+
if use_dynamic_roi:
2812+
return _emit_resize_topi_dynamic_roi(
2813+
bb,
2814+
x,
2815+
roi_dynamic_vec,
2816+
sizes,
2817+
ndims,
2818+
topi_mode,
2819+
coord_mode,
2820+
rounding_method,
2821+
cubic_coeff_a,
2822+
exclude_outside,
2823+
extrapolation_value,
2824+
)
27082825

27092826
if ndims == 3:
27102827
return bb.emit_te(
27112828
topi.image.resize1d,
27122829
x,
2713-
roi,
2830+
roi_static,
27142831
sizes,
27152832
"NCW",
27162833
topi_mode,
@@ -2724,7 +2841,7 @@ def _impl_v18(cls, bb, inputs, attr, params):
27242841
return relax.op.image.resize2d(
27252842
x,
27262843
size=relax.ShapeExpr(sizes),
2727-
roi=roi,
2844+
roi=roi_static,
27282845
layout="NCHW",
27292846
method=relax_mode,
27302847
coordinate_transformation_mode=coord_mode,
@@ -2734,10 +2851,11 @@ def _impl_v18(cls, bb, inputs, attr, params):
27342851
extrapolation_value=extrapolation_value,
27352852
)
27362853
else: # ndims == 5
2854+
roi3d = _topi_resize3d_roi_from_onnx_ncdhw_spatial(roi_static)
27372855
return relax.op.image.resize3d(
27382856
x,
27392857
size=relax.ShapeExpr(sizes),
2740-
roi=roi,
2858+
roi=roi3d,
27412859
layout="NCDHW",
27422860
method=relax_mode,
27432861
coordinate_transformation_mode=coord_mode,

tests/python/relax/test_frontend_onnx.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3268,6 +3268,64 @@ def test_resize(with_roi, roi_list, with_constant):
32683268
check_correctness(model)
32693269

32703270

3271+
def test_resize_dynamic_roi_tf_crop_and_resize():
3272+
"""ROI is a graph input (not initializer), lowered through TOPI dynamic-ROI path."""
3273+
resize_node = helper.make_node(
3274+
"Resize",
3275+
["X", "roi", "scales"],
3276+
["Y"],
3277+
mode="linear",
3278+
coordinate_transformation_mode="tf_crop_and_resize",
3279+
)
3280+
graph = helper.make_graph(
3281+
[resize_node],
3282+
"resize_dynamic_roi",
3283+
inputs=[
3284+
helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 32, 32]),
3285+
helper.make_tensor_value_info("roi", TensorProto.FLOAT, [8]),
3286+
],
3287+
initializer=[
3288+
helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0]),
3289+
],
3290+
outputs=[
3291+
helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 64, 64]),
3292+
],
3293+
)
3294+
model = helper.make_model(graph, producer_name="resize_dynamic_roi")
3295+
check_correctness(model, atol=1e-5)
3296+
3297+
3298+
def test_resize_dynamic_roi_3d_tf_crop_and_resize():
3299+
"""5-D NCDHW: ROI is a graph input; covers dynamic-ROI TOPI resize3d path."""
3300+
resize_node = helper.make_node(
3301+
"Resize",
3302+
["X", "roi", "scales"],
3303+
["Y"],
3304+
mode="linear",
3305+
coordinate_transformation_mode="tf_crop_and_resize",
3306+
)
3307+
graph = helper.make_graph(
3308+
[resize_node],
3309+
"resize_dynamic_roi_3d",
3310+
inputs=[
3311+
helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 4, 5]),
3312+
helper.make_tensor_value_info("roi", TensorProto.FLOAT, [10]),
3313+
],
3314+
initializer=[
3315+
helper.make_tensor("scales", TensorProto.FLOAT, [5], [1.0, 1.0, 2.0, 2.0, 2.0]),
3316+
],
3317+
outputs=[
3318+
helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 6, 8, 10]),
3319+
],
3320+
)
3321+
model = helper.make_model(graph, producer_name="resize_dynamic_roi_3d")
3322+
# Use a valid full-tensor ROI so ORT and TOPI agree on tf_crop_and_resize (random ROI
3323+
# can hit extrapolation / numerical differences across runtimes).
3324+
x_np = rg.standard_normal((1, 1, 3, 4, 5)).astype(np.float32)
3325+
roi_np = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], dtype=np.float32)
3326+
check_correctness(model, opset=18, atol=1e-5, inputs={"X": x_np, "roi": roi_np})
3327+
3328+
32713329
def test_resize_nd_sizes():
32723330
cases = [
32733331
("resize1d", [1, 1, 4], [1, 1, 7]),

0 commit comments

Comments
 (0)