From b7e0cfb6d469c3745ae2195908daadea9c64d87e Mon Sep 17 00:00:00 2001 From: masahi Date: Thu, 18 Feb 2021 12:11:17 +0900 Subject: [PATCH] [TOPI, Relay] Support roi_align NHWC layout (#7463) * begin nhwc roi align * integrate mode change from upstream * adding test * support nhwc shape func * update strategy * refactoring test * refactor test * refactoring * fix lint * update relay op tests --- python/tvm/relay/op/strategy/cuda.py | 20 +- python/tvm/relay/op/strategy/generic.py | 20 +- python/tvm/relay/op/strategy/x86.py | 19 +- python/tvm/relay/op/vision/_vision.py | 17 +- python/tvm/topi/testing/__init__.py | 2 +- python/tvm/topi/testing/roi_align_python.py | 153 +++++++++++---- python/tvm/topi/vision/rcnn/roi_align.py | 196 +++++++++++++++----- tests/python/relay/test_op_level5.py | 89 +++++++-- 8 files changed, 398 insertions(+), 118 deletions(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 032d2dd2c8f1..cb4688c4889e 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -945,12 +945,20 @@ def roi_align_strategy_cuda(attrs, inputs, out_type, target): """roi_align cuda strategy""" strategy = _op.OpStrategy() layout = attrs.layout - assert layout == "NCHW", "only support nchw for now" - strategy.add_implementation( - wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw), - wrap_topi_schedule(topi.cuda.schedule_roi_align), - name="roi_align_nchw.cuda", - ) + + if layout == "NCHW": + strategy.add_implementation( + wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw), + wrap_topi_schedule(topi.cuda.schedule_roi_align), + name="roi_align_nchw.cuda", + ) + else: + assert layout == "NHWC", "layout must be NCHW or NHWC." + strategy.add_implementation( + wrap_compute_roi_align(topi.vision.rcnn.roi_align_nhwc), + wrap_topi_schedule(topi.cuda.schedule_roi_align), + name="roi_align_nhwc.cuda", + ) return strategy diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index e744b8c9da83..f076176c5d8a 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -1039,7 +1039,6 @@ def wrap_compute_roi_align(topi_compute): """wrap roi_align topi compute""" def _compute_roi_align(attrs, inputs, out_type): - assert attrs.layout == "NCHW" pooled_size = get_const_tuple(attrs.pooled_size) mode = bytes(attrs.mode, "utf-8") return [ @@ -1061,12 +1060,19 @@ def roi_align_strategy(attrs, inputs, out_type, target): """roi_align generic strategy""" strategy = _op.OpStrategy() layout = attrs.layout - assert layout == "NCHW", "only support nchw for now" - strategy.add_implementation( - wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw), - wrap_topi_schedule(topi.generic.schedule_roi_align), - name="roi_align.generic", - ) + if layout == "NCHW": + strategy.add_implementation( + wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw), + wrap_topi_schedule(topi.generic.schedule_roi_align), + name="roi_align.generic", + ) + else: + assert layout == "NHWC", "layout must be NCHW or NHWC." + strategy.add_implementation( + wrap_compute_roi_align(topi.vision.rcnn.roi_align_nhwc), + wrap_topi_schedule(topi.generic.schedule_roi_align), + name="roi_align.generic", + ) return strategy diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py index f33c45b248d6..1f37a4f8e98c 100644 --- a/python/tvm/relay/op/strategy/x86.py +++ b/python/tvm/relay/op/strategy/x86.py @@ -481,12 +481,19 @@ def roi_align_strategy_cpu(attrs, inputs, out_type, target): """roi_align x86 strategy""" strategy = _op.OpStrategy() layout = attrs.layout - assert layout == "NCHW", "only support nchw for now" - strategy.add_implementation( - wrap_compute_roi_align(topi.x86.roi_align_nchw), - wrap_topi_schedule(topi.generic.schedule_roi_align), - name="roi_align.x86", - ) + if layout == "NCHW": + strategy.add_implementation( + wrap_compute_roi_align(topi.x86.roi_align_nchw), + wrap_topi_schedule(topi.generic.schedule_roi_align), + name="roi_align.x86", + ) + else: + assert layout == "NHWC", "layout must be NCHW or NHWC." + strategy.add_implementation( + wrap_compute_roi_align(topi.vision.rcnn.roi_align_nhwc), + wrap_topi_schedule(topi.generic.schedule_roi_align), + name="roi_align.x86", + ) return strategy diff --git a/python/tvm/relay/op/vision/_vision.py b/python/tvm/relay/op/vision/_vision.py index 04676e24adf6..9c8c853fa3d2 100644 --- a/python/tvm/relay/op/vision/_vision.py +++ b/python/tvm/relay/op/vision/_vision.py @@ -86,7 +86,7 @@ def nms_shape_func(attrs, inputs, _): @script -def _roi_align_shape_func(data_shape, rois_shape, pooled_size): +def _roi_align_shape_func_nchw(data_shape, rois_shape, pooled_size): out = output_tensor((4,), "int64") out[0] = rois_shape[0] out[1] = data_shape[1] @@ -95,6 +95,19 @@ def _roi_align_shape_func(data_shape, rois_shape, pooled_size): return out +@script +def _roi_align_shape_func_nhwc(data_shape, rois_shape, pooled_size): + out = output_tensor((4,), "int64") + out[0] = rois_shape[0] + out[1] = int64(pooled_size[0]) + out[2] = int64(pooled_size[1]) + out[3] = data_shape[3] + return out + + @reg.register_shape_func("vision.roi_align", False) def roi_align_shape_func(attrs, inputs, _): - return [_roi_align_shape_func(inputs[0], inputs[1], convert(attrs.pooled_size))] + if attrs.layout == "NCHW": + return [_roi_align_shape_func_nchw(inputs[0], inputs[1], convert(attrs.pooled_size))] + assert attrs.layout == "NHWC", "layout must be NCHW or NHWC." + return [_roi_align_shape_func_nhwc(inputs[0], inputs[1], convert(attrs.pooled_size))] diff --git a/python/tvm/topi/testing/__init__.py b/python/tvm/topi/testing/__init__.py index 85f13a763c40..ef36b9e73446 100644 --- a/python/tvm/topi/testing/__init__.py +++ b/python/tvm/topi/testing/__init__.py @@ -39,7 +39,7 @@ from .bilinear_resize_python import bilinear_resize_python from .trilinear_resize3d_python import trilinear_resize3d_python from .reorg_python import reorg_python -from .roi_align_python import roi_align_nchw_python +from .roi_align_python import roi_align_nchw_python, roi_align_nhwc_python from .roi_pool_python import roi_pool_nchw_python from .lrn_python import lrn_python from .l2_normalize_python import l2_normalize_python diff --git a/python/tvm/topi/testing/roi_align_python.py b/python/tvm/topi/testing/roi_align_python.py index 643a954b101b..986123b6c9c6 100644 --- a/python/tvm/topi/testing/roi_align_python.py +++ b/python/tvm/topi/testing/roi_align_python.py @@ -20,42 +20,51 @@ import numpy as np -def roi_align_nchw_python(a_np, rois_np, pooled_size, spatial_scale, sample_ratio, mode=b"avg"): - """Roi align in python""" - avg_mode = mode in (b"avg", "avg", 0) - max_mode = mode in (b"max", "max", 1) - assert avg_mode or max_mode, "Mode must be average or max. Please pass a valid mode." - _, channel, height, width = a_np.shape - num_roi = rois_np.shape[0] - b_np = np.zeros((num_roi, channel, pooled_size, pooled_size), dtype=a_np.dtype) - if isinstance(pooled_size, int): - pooled_size_h = pooled_size_w = pooled_size - else: - pooled_size_h, pooled_size_w = pooled_size - - def _bilinear(n, c, y, x): - if y < -1 or y > height or x < -1 or x > width: - return 0 +def _bilinear(a_np, n, c, y, x, height, width, layout): + if y < -1 or y > height or x < -1 or x > width: + return 0 - y = min(max(y, 0), height - 1) - x = min(max(x, 0), width - 1) + y = min(max(y, 0), height - 1) + x = min(max(x, 0), width - 1) - y_low = int(math.floor(y)) - x_low = int(math.floor(x)) - y_high = y_low + 1 - x_high = x_low + 1 + y_low = int(math.floor(y)) + x_low = int(math.floor(x)) + y_high = y_low + 1 + x_high = x_low + 1 - wy_h = y - y_low - wx_h = x - x_low - wy_l = 1 - wy_h - wx_l = 1 - wx_h + wy_h = y - y_low + wx_h = x - x_low + wy_l = 1 - wy_h + wx_l = 1 - wx_h - val = 0 - for wx, xp in zip((wx_l, wx_h), (x_low, x_high)): - for wy, yp in zip((wy_l, wy_h), (y_low, y_high)): - if 0 <= yp < height and 0 <= xp < width: + val = 0 + for wx, xp in zip((wx_l, wx_h), (x_low, x_high)): + for wy, yp in zip((wy_l, wy_h), (y_low, y_high)): + if 0 <= yp < height and 0 <= xp < width: + if layout == "NCHW": val += wx * wy * a_np[n, c, yp, xp] - return val + else: + val += wx * wy * a_np[n, yp, xp, c] + return val + + +def roi_align_common( + a_np, + b_np, + rois_np, + channel, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + avg_mode, + max_mode, + height, + width, + layout, +): + """Common code used by roi align NCHW and NHWC""" + num_roi = rois_np.shape[0] for i in range(num_roi): roi = rois_np[i] @@ -70,8 +79,8 @@ def _bilinear(n, c, y, x): if sample_ratio > 0: roi_bin_grid_h = roi_bin_grid_w = int(sample_ratio) else: - roi_bin_grid_h = int(math.ceil(roi_h / pooled_size)) - roi_bin_grid_w = int(math.ceil(roi_w / pooled_size)) + roi_bin_grid_h = int(math.ceil(roi_h / pooled_size_h)) + roi_bin_grid_w = int(math.ceil(roi_w / pooled_size_w)) count = roi_bin_grid_h * roi_bin_grid_w @@ -87,8 +96,80 @@ def _bilinear(n, c, y, x): y = roi_start_h + ph * bin_h + (iy + 0.5) * bin_h / roi_bin_grid_h x = roi_start_w + pw * bin_w + (ix + 0.5) * bin_w / roi_bin_grid_w if avg_mode: - total += _bilinear(batch_index, c, y, x) / count + total += ( + _bilinear(a_np, batch_index, c, y, x, height, width, layout) + / count + ) if max_mode: - total = max(total, _bilinear(batch_index, c, y, x)) - b_np[i, c, ph, pw] = total + total = max( + total, + _bilinear(a_np, batch_index, c, y, x, height, width, layout), + ) + + if layout == "NCHW": + b_np[i, c, ph, pw] = total + else: + b_np[i, ph, pw, c] = total return b_np + + +def roi_align_nchw_python(a_np, rois_np, pooled_size, spatial_scale, sample_ratio, mode=b"avg"): + """Roi align NCHW in python""" + avg_mode = mode in (b"avg", "avg", 0) + max_mode = mode in (b"max", "max", 1) + assert avg_mode or max_mode, "Mode must be average or max. Please pass a valid mode." + _, channel, height, width = a_np.shape + if isinstance(pooled_size, int): + pooled_size_h = pooled_size_w = pooled_size + else: + pooled_size_h, pooled_size_w = pooled_size + + b_np = np.zeros((rois_np.shape[0], channel, pooled_size_h, pooled_size_w), dtype=a_np.dtype) + + return roi_align_common( + a_np, + b_np, + rois_np, + channel, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + avg_mode, + max_mode, + height, + width, + "NCHW", + ) + + +def roi_align_nhwc_python(a_np, rois_np, pooled_size, spatial_scale, sample_ratio, mode=b"avg"): + """Roi align NHWC in python""" + avg_mode = mode in (b"avg", "avg", 0) + max_mode = mode in (b"max", "max", 1) + assert avg_mode or max_mode, "Mode must be average or max. Please pass a valid mode." + _, height, width, channel = a_np.shape + num_roi = rois_np.shape[0] + + if isinstance(pooled_size, int): + pooled_size_h = pooled_size_w = pooled_size + else: + pooled_size_h, pooled_size_w = pooled_size + + b_np = np.zeros((num_roi, pooled_size_h, pooled_size_w, channel), dtype=a_np.dtype) + + return roi_align_common( + a_np, + b_np, + rois_np, + channel, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + avg_mode, + max_mode, + height, + width, + "NHWC", + ) diff --git a/python/tvm/topi/vision/rcnn/roi_align.py b/python/tvm/topi/vision/rcnn/roi_align.py index 95f350084ba5..655ba2637d84 100644 --- a/python/tvm/topi/vision/rcnn/roi_align.py +++ b/python/tvm/topi/vision/rcnn/roi_align.py @@ -19,7 +19,71 @@ import tvm from tvm import te from ...utils import get_const_tuple -from ...cpp.utils import bilinear_sample_nchw +from ...cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc + + +def _sample_common( + i, + c, + ph, + pw, + rois, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + dtype, + avg_mode, + bilinear_func, +): + roi = rois[i] + batch_index = roi[0].astype("int32") + roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1], roi[2], roi[3], roi[4] + roi_start_h *= spatial_scale + roi_end_h *= spatial_scale + roi_start_w *= spatial_scale + roi_end_w *= spatial_scale + + # force malformed ROIs to be 1x1 + roi_h = tvm.te.max(roi_end_h - roi_start_h, tvm.tir.const(1.0, dtype)) + roi_w = tvm.te.max(roi_end_w - roi_start_w, tvm.tir.const(1.0, dtype)) + + bin_h = roi_h / pooled_size_h + bin_w = roi_w / pooled_size_w + + if sample_ratio > 0: + roi_bin_grid_h = roi_bin_grid_w = tvm.tir.const(sample_ratio, "int32") + else: + roi_bin_grid_h = te.ceil(roi_h / pooled_size_h).astype("int32") + roi_bin_grid_w = te.ceil(roi_w / pooled_size_w).astype("int32") + + count = roi_bin_grid_h * roi_bin_grid_w + rh = te.reduce_axis((0, roi_bin_grid_h)) + rw = te.reduce_axis((0, roi_bin_grid_w)) + roi_start_h += ph * bin_h + roi_start_w += pw * bin_w + + if avg_mode: + return te.sum( + bilinear_func( + batch_index, + c, + roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h, + roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w, + ) + / count, + axis=[rh, rw], + ) + # max mode + return te.max( + bilinear_func( + batch_index, + c, + roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h, + roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w, + ), + axis=[rh, rw], + ) def roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1): @@ -73,54 +137,92 @@ def _bilinear(i, c, y, x): return tvm.tir.if_then_else(outside, 0.0, val) def _sample(i, c, ph, pw): - roi = rois[i] - batch_index = roi[0].astype("int32") - roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1], roi[2], roi[3], roi[4] - roi_start_h *= spatial_scale - roi_end_h *= spatial_scale - roi_start_w *= spatial_scale - roi_end_w *= spatial_scale - - # force malformed ROIs to be 1x1 - roi_h = tvm.te.max(roi_end_h - roi_start_h, tvm.tir.const(1.0, dtype)) - roi_w = tvm.te.max(roi_end_w - roi_start_w, tvm.tir.const(1.0, dtype)) - - bin_h = roi_h / pooled_size_h - bin_w = roi_w / pooled_size_w - - if sample_ratio > 0: - roi_bin_grid_h = roi_bin_grid_w = tvm.tir.const(sample_ratio, "int32") - else: - roi_bin_grid_h = te.ceil(roi_h / pooled_size_h).astype("int32") - roi_bin_grid_w = te.ceil(roi_w / pooled_size_w).astype("int32") - - count = roi_bin_grid_h * roi_bin_grid_w - rh = te.reduce_axis((0, roi_bin_grid_h)) - rw = te.reduce_axis((0, roi_bin_grid_w)) - roi_start_h += ph * bin_h - roi_start_w += pw * bin_w - if avg_mode: - return te.sum( - _bilinear( - batch_index, - c, - roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h, - roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w, - ) - / count, - axis=[rh, rw], - ) - # max mode - return te.max( - _bilinear( - batch_index, - c, - roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h, - roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w, - ), - axis=[rh, rw], + return _sample_common( + i, + c, + ph, + pw, + rois, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + dtype, + avg_mode, + _bilinear, ) return te.compute( (num_roi, channel, pooled_size_h, pooled_size_w), _sample, tag="pool,roi_align_nchw" ) + + +def roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1): + """ROI align operator in NHWC layout. + + Parameters + ---------- + data : tvm.te.Tensor + 4-D with shape [batch, height, width, channel] + + rois : tvm.te.Tensor + 2-D with shape [num_roi, 5]. The last dimension should be in format of + [batch_index, w_start, h_start, w_end, h_end] + + pooled_size : int or list/tuple of two ints + output size, or [out_height, out_width] + + spatial_scale : float + Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal + of total stride in convolutional layers, which should be in range (0.0, 1.0] + + mode : int or str + There are two modes, average and max. For the average mode, you can pass b'avg' or 0, and + for the max mode, you can pass b'max' or 1. + + sample_ratio : int + Optional sampling ratio of ROI align, using adaptive size by default. + + Returns + ------- + output : tvm.te.Tensor + 4-D with shape [num_roi, pooled_size, pooled_size, channel] + """ + avg_mode = mode in (b"avg", 0) + max_mode = mode in (b"max", 1) + assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode." + dtype = rois.dtype + _, height, width, channel = get_const_tuple(data.shape) + num_roi, _ = get_const_tuple(rois.shape) + + if isinstance(pooled_size, int): + pooled_size_h = pooled_size_w = pooled_size + else: + pooled_size_h, pooled_size_w = pooled_size + + def _bilinear(i, c, y, x): + outside = tvm.tir.any(y < -1.0, x < -1.0, y > height, x > width) + y = tvm.te.min(tvm.te.max(y, 0.0), height - 1) + x = tvm.te.min(tvm.te.max(x, 0.0), width - 1) + val = bilinear_sample_nhwc(data, (i, y, x, c), height - 1, width - 1) + return tvm.tir.if_then_else(outside, 0.0, val) + + def _sample(i, ph, pw, c): + return _sample_common( + i, + c, + ph, + pw, + rois, + pooled_size_h, + pooled_size_w, + spatial_scale, + sample_ratio, + dtype, + avg_mode, + _bilinear, + ) + + return te.compute( + (num_roi, pooled_size_h, pooled_size_w, channel), _sample, tag="pool,roi_align_nchw" + ) diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index 95cd537091f5..0a84667f8bdb 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -583,7 +583,18 @@ def test_threshold(): @tvm.testing.uses_gpu def test_roi_align(): - def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ratio, mode): + def verify_roi_align( + data_shape, + rois_shape, + channel, + in_size, + pooled_size, + spatial_scale, + sample_ratio, + mode, + layout, + ref_func, + ): data = relay.var("data", relay.ty.TensorType(data_shape, "float32")) rois = relay.var("rois", relay.ty.TensorType(rois_shape, "float32")) z = relay.vision.roi_align( @@ -593,21 +604,27 @@ def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ spatial_scale=spatial_scale, sample_ratio=sample_ratio, mode=mode, - layout="NCHW", + layout=layout, ) zz = run_infer_type(z) - batch, channel, in_size, _ = data_shape + num_roi = rois_shape[0] - assert zz.checked_type == relay.ty.TensorType( - (num_roi, channel, pooled_size, pooled_size), "float32" - ) + + if layout == "NCHW": + assert zz.checked_type == relay.ty.TensorType( + (num_roi, channel, pooled_size, pooled_size), "float32" + ) + else: + assert zz.checked_type == relay.ty.TensorType( + (num_roi, pooled_size, pooled_size, channel), "float32" + ) func = relay.Function([data, rois], z) func = run_infer_type(func) np_data = np.random.uniform(size=data_shape).astype("float32") np_rois = np.random.uniform(size=rois_shape).astype("float32") * in_size - np_rois[:, 0] = np.random.randint(low=0, high=batch, size=num_roi) - ref_res = tvm.topi.testing.roi_align_nchw_python( + np_rois[:, 0] = np.random.randint(low=0, high=data_shape[0], size=num_roi) + ref_res = ref_func( np_data, np_rois, pooled_size=pooled_size, @@ -616,6 +633,7 @@ def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ mode=mode, ) for target, ctx in tvm.testing.enabled_targets(): + print("test on", target) intrp1 = relay.create_executor("graph", ctx=ctx, target=target) op_res1 = intrp1.evaluate(func)(np_data, np_rois) tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-4) @@ -623,18 +641,64 @@ def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ op_res2 = intrp2.evaluate(func)(np_data, np_rois) tvm.testing.assert_allclose(op_res2.asnumpy(), ref_res, rtol=1e-4) - verify_roi_align( + def verify_roi_align_nchw( + data_shape, rois_shape, pooled_size, spatial_scale, sample_ratio, mode + ): + _, channel, in_size, _ = data_shape + return verify_roi_align( + data_shape, + rois_shape, + channel, + in_size, + pooled_size, + spatial_scale, + sample_ratio, + mode, + "NCHW", + tvm.topi.testing.roi_align_nchw_python, + ) + + def verify_roi_align_nhwc( + data_shape, rois_shape, pooled_size, spatial_scale, sample_ratio, mode + ): + _, in_size, _, channel = data_shape + return verify_roi_align( + data_shape, + rois_shape, + channel, + in_size, + pooled_size, + spatial_scale, + sample_ratio, + mode, + "NHWC", + tvm.topi.testing.roi_align_nhwc_python, + ) + + verify_roi_align_nchw( (1, 4, 16, 16), (32, 5), pooled_size=7, spatial_scale=1.0, sample_ratio=-1, mode="avg" ) - verify_roi_align( + verify_roi_align_nchw( (4, 4, 16, 16), (32, 5), pooled_size=7, spatial_scale=0.5, sample_ratio=2, mode="avg" ) - verify_roi_align( + verify_roi_align_nchw( (1, 4, 16, 16), (32, 5), pooled_size=7, spatial_scale=1.0, sample_ratio=-1, mode="max" ) - verify_roi_align( + verify_roi_align_nchw( (4, 4, 16, 16), (32, 5), pooled_size=7, spatial_scale=0.5, sample_ratio=2, mode="max" ) + verify_roi_align_nhwc( + (1, 16, 16, 4), (32, 5), pooled_size=7, spatial_scale=1.0, sample_ratio=-1, mode="avg" + ) + verify_roi_align_nhwc( + (4, 16, 16, 4), (32, 5), pooled_size=7, spatial_scale=0.5, sample_ratio=2, mode="avg" + ) + verify_roi_align_nhwc( + (1, 16, 16, 4), (32, 5), pooled_size=7, spatial_scale=1.0, sample_ratio=-1, mode="max" + ) + verify_roi_align_nhwc( + (4, 16, 16, 4), (32, 5), pooled_size=7, spatial_scale=0.5, sample_ratio=2, mode="max" + ) @tvm.testing.uses_gpu @@ -1262,7 +1326,6 @@ def verify_batch_to_space_nd(dshape, block_shape, crops): test_resize_infer_type() test_resize() test_resize3d_infer_type() - test_resize3d() test_crop_and_resize() test_multibox_prior() test_multibox_transform_loc()