diff --git a/src/relay/op/nn/pooling.cc b/src/relay/op/nn/pooling.cc index dfda0889b6c3..8f1115d56c83 100644 --- a/src/relay/op/nn/pooling.cc +++ b/src/relay/op/nn/pooling.cc @@ -947,21 +947,21 @@ bool Pool3DRel(const Array& types, const auto hidx = layout.IndexOf(LayoutAxis::Get('H')); const auto widx = layout.IndexOf(LayoutAxis::Get('W')); - IndexExpr pad_d, pad_h, pad_w; + IndexExpr pad[3]; if (param->padding.size() == 1) { - pad_d = param->padding[0] * 2; - pad_h = param->padding[0] * 2; - pad_w = param->padding[0] * 2; + pad[0] = param->padding[0] * 2; + pad[1] = param->padding[0] * 2; + pad[2] = param->padding[0] * 2; } else if (param->padding.size() == 3) { // (front, top, left) - pad_d = param->padding[0] * 2; - pad_h = param->padding[1] * 2; - pad_w = param->padding[2] * 2; + pad[0] = param->padding[0] * 2; + pad[1] = param->padding[1] * 2; + pad[2] = param->padding[2] * 2; } else if (param->padding.size() == 6) { // (front, top, left, back, bottom, right) - pad_d = param->padding[0] + param->padding[3]; - pad_h = param->padding[1] + param->padding[4]; - pad_w = param->padding[2] + param->padding[5]; + pad[0] = param->padding[0] + param->padding[3]; + pad[1] = param->padding[1] + param->padding[4]; + pad[2] = param->padding[2] + param->padding[5]; } else { return false; } @@ -978,10 +978,10 @@ bool Pool3DRel(const Array& types, oshape[ii] = dshape[ii]; } else { if (param->ceil_mode) { - oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i] + + oshape[ii] = ((dshape[ii] + pad[i] - param->pool_size[i] + param->strides[i] - 1) / param->strides[i]) + 1; } else { - oshape[ii] = ((dshape[ii] + pad_d - param->pool_size[i]) / param->strides[i]) + 1; + oshape[ii] = ((dshape[ii] + pad[i] - param->pool_size[i]) / param->strides[i]) + 1; } } } diff --git a/tests/python/relay/test_op_level2.py b/tests/python/relay/test_op_level2.py index 27eab0eeb9b0..0acd83639363 100644 --- a/tests/python/relay/test_op_level2.py +++ b/tests/python/relay/test_op_level2.py @@ -21,16 +21,10 @@ from tvm import autotvm from tvm import relay from tvm.relay import transform -from tvm.relay.testing import ctx_list +from tvm.relay.testing import ctx_list, run_infer_type from tvm.contrib import util import topi.testing -def run_infer_type(expr): - mod = relay.Module.from_expr(expr) - mod = transform.InferType()(mod) - entry = mod["main"] - return entry if isinstance(expr, relay.Function) else entry.body - def test_conv1d_infer_type(): # symbolic in batch dimension @@ -768,7 +762,7 @@ def _test_pool1d(opfunc): def test_pool3d(): - def _test_pool3d(opfunc): + def _test_pool3d(opfunc, padding=(0, 0, 0, 0, 0, 0), out_shape=(1, 3, 16, 16, 16)): n, c, d, h, w = tvm.size_var("n"), 10, 5, 224, 224 x = relay.var("x", relay.TensorType((n, c, d, h, w), "float32")) y = opfunc(x, pool_size=(1, 1, 1)) @@ -780,18 +774,28 @@ def _test_pool3d(opfunc): dshape = (1, 3, 32, 32, 32) x = relay.var("x", shape=dshape) pool_type = 'max' if 'max' in str(opfunc) else 'avg' - y = opfunc(x, pool_size=(2, 2, 2), strides=(2, 2, 2), padding=(0, 0, 0, 0, 0, 0)) + y = opfunc(x, pool_size=(2, 2, 2), strides=(2, 2, 2), padding=padding) func = relay.Function([x], y) + # check output shape + f_out_shape = tuple(map(lambda x: int(x), run_infer_type(func).ret_type.shape)) + assert out_shape == f_out_shape, \ + "Output shape mismatch. expected {}, actual {}".format(out_shape, f_out_shape) data = np.random.uniform(size=dshape).astype(dtype) ref_res = topi.testing.pool3d_ncdhw_python(data, (2, 2, 2), (2, 2, 2), - (0, 0, 0, 0, 0, 0), (1, 3, 16, 16, 16), pool_type, False) + padding, out_shape, pool_type, False) for target, ctx in ctx_list(): intrp1 = relay.create_executor("graph", ctx=ctx, target=target) op_res1 = intrp1.evaluate(func)(data) tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5) _test_pool3d(relay.nn.max_pool3d) + _test_pool3d(relay.nn.max_pool3d, padding=(2, 0, 0, 2, 0, 0), out_shape=(1, 3, 18, 16, 16)) + _test_pool3d(relay.nn.max_pool3d, padding=(0, 3, 0, 0, 3, 0), out_shape=(1, 3, 16, 19, 16)) + _test_pool3d(relay.nn.max_pool3d, padding=(0, 0, 4, 0, 0, 4), out_shape=(1, 3, 16, 16, 20)) _test_pool3d(relay.nn.avg_pool3d) + _test_pool3d(relay.nn.avg_pool3d, padding=(2, 0, 0, 2, 0, 0), out_shape=(1, 3, 18, 16, 16)) + _test_pool3d(relay.nn.avg_pool3d, padding=(0, 3, 0, 0, 3, 0), out_shape=(1, 3, 16, 19, 16)) + _test_pool3d(relay.nn.avg_pool3d, padding=(0, 0, 4, 0, 0, 4), out_shape=(1, 3, 16, 16, 20)) def test_avg_pool2d_no_count_pad(): diff --git a/topi/python/topi/testing/pool3d_python.py b/topi/python/topi/testing/pool3d_python.py index 32513163d068..631a995e7c12 100644 --- a/topi/python/topi/testing/pool3d_python.py +++ b/topi/python/topi/testing/pool3d_python.py @@ -18,6 +18,7 @@ """max_pool3d and avg_pool3d in python""" import math import numpy as np +import tvm def pool3d_ncdhw_python(np_data, kernel, strides, padding, @@ -39,10 +40,17 @@ def pool3d_ncdhw_python(np_data, kernel, assert out_shape[3] == int(math.floor(float(in_shape[3] - k_h + pt + pb) / s_h) + 1) assert out_shape[4] == int(math.floor(float(in_shape[4] - k_w + pl + pr) / s_w) + 1) - pad_np = np.zeros(shape=(in_n, in_c, - in_d + pf + pk, - in_h + pt + pb, - in_w + pl + pr)).astype(dtype) + fill_value = tvm.const(0.0, dtype).value + if not(count_include_pad) and pool_type == 'max': + fill_value = tvm.min_value(dtype).value + + pad_np = np.full(shape=(in_n, in_c, + in_d + pf + pk, + in_h + pt + pb, + in_w + pl + pr), + fill_value=fill_value, + dtype=dtype) + no_zero = (range(in_n), range(in_c), (range(pf, in_d + pf)), @@ -81,5 +89,5 @@ def pool3d_ncdhw_python(np_data, kernel, else: raise ValueError("pool type {} is not supported".format(pool_type)) - ret_np = np.maximum(ret_np, 0.0) + ret_np = np.maximum(ret_np, fill_value) return ret_np