Fix padding in pooling op (#4738)

src/relay/op/nn/pooling.cc

@@ -947,21 +947,21 @@ bool Pool3DRel(const Array<Type>& 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<Type>& 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;
       }
     }
   }

tests/python/relay/test_op_level2.py

@@ -21,16 +21,10 @@ import tvm
 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():
 
 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():
         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():

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