add more tests including NDHWC

tests/python/relay/test_op_level10.py

@@ -356,7 +356,7 @@ def verify_adaptive_pool(dshape, out_size, pool_type, layout, dtype, opfunc):
     func = relay.Function([x], y)
 
     np_data = np.random.uniform(low=0, high=255, size=dshape).astype(dtype)
-    np_out = topi.testing.adaptive_pool(np_data, out_size, pool_type)
+    np_out = topi.testing.adaptive_pool(np_data, out_size, pool_type, layout)
 
     for target, ctx in ctx_list():
         intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
@@ -379,8 +379,12 @@ def test_adaptive_pool():
     verify_adaptive_pool2d((1, 3, 224, 224), (2, 3), "avg")
     verify_adaptive_pool2d((1, 14, 56, 78), (34, 13), "max")
     verify_adaptive_pool2d((1, 5, 46, 97), (4, 96), "avg")
+    verify_adaptive_pool2d((1, 224, 224, 3), (1, 1), "max", layout="NHWC")
+    verify_adaptive_pool2d((1, 3, 224, 224), (2, 3), "avg", layout="NHWC")
     verify_adaptive_pool3d((1, 16, 32, 32, 32), (1, 1, 1), "max", layout="NCDHW")
     verify_adaptive_pool3d((1, 16, 32, 32, 32), (1, 1, 1), "avg", layout="NCDHW")
+    verify_adaptive_pool3d((1, 16, 32, 32, 32), (1, 1, 1), "avg", layout="NDHWC")
+    verify_adaptive_pool3d((1, 16, 32, 32, 32), (2, 4, 4), "max", layout="NDHWC")
 
 
 def test_sequence_mask():

topi/python/topi/testing/adaptive_pool_python.py

@@ -19,59 +19,93 @@
 import numpy as np
 
 
-def adaptive_pool(np_data, out_size, pool_type):
-    """ The reference function for adaptive pool, both 2d and 3d """
-    def start_index(index, odim, idim):
-        return int(np.floor(index * idim / odim))
-
-    def end_index(index, odim, idim):
-        return int(np.ceil((index + 1) * idim / odim))
-
-    def pool2d(i, j):
-        out = np.zeros(out_size).astype(np_data.dtype)
-        n, c, h, w = np_data.shape
-        oh, ow = out_size
+def _start_index(index, odim, idim):
+    return int(np.floor(index * idim / odim))
+
+
+def _end_index(index, odim, idim):
+    return int(np.ceil((index + 1) * idim / odim))
+
+
+def _pool2d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    oh, ow = out_size
+    for k in range(oh):
+        k_start = _start_index(k, oh, in_size[0])
+        k_end = _end_index(k, oh, in_size[0])
+        k_sl = slice(k_start, k_end)
+        for l in range(ow):
+            l_start = _start_index(l, ow, in_size[1])
+            l_end = _end_index(l, ow, in_size[1])
+            l_sl = slice(l_start, l_end)
+            out[k, l] = np_op(np_data[k_sl, l_sl])
+    return out
+
+
+def _pool3d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    od, oh, ow = out_size
+    for m in range(od):
+        m_start = _start_index(m, od, in_size[0])
+        m_end = _end_index(m, od, in_size[0])
+        m_sl = slice(m_start, m_end)
         for k in range(oh):
-            k_start = start_index(k, oh, h)
-            k_end = end_index(k, oh, h)
+            k_start = _start_index(k, oh, in_size[1])
+            k_end = _end_index(k, oh, in_size[1])
             k_sl = slice(k_start, k_end)
             for l in range(ow):
-                l_start = start_index(l, ow, w)
-                l_end = end_index(l, ow, w)
+                l_start = _start_index(l, ow, in_size[2])
+                l_end = _end_index(l, ow, in_size[2])
                 l_sl = slice(l_start, l_end)
-                out[k, l] = np_op(np_data[i, j, k_sl, l_sl])
-        return out
-
-    def pool3d(i, j):
-        out = np.zeros(out_size).astype(np_data.dtype)
-        n, c, d, h, w = np_data.shape
-        od, oh, ow = out_size
-        for m in range(od):
-            m_start = start_index(m, od, d)
-            m_end = end_index(m, od, d)
-            m_sl = slice(m_start, m_end)
-            for k in range(oh):
-                k_start = start_index(k, oh, h)
-                k_end = end_index(k, oh, h)
-                k_sl = slice(k_start, k_end)
-                for l in range(ow):
-                    l_start = start_index(l, ow, w)
-                    l_end = end_index(l, ow, w)
-                    l_sl = slice(l_start, l_end)
-                    out[m, k, l] = np_op(np_data[i, j, m_sl, k_sl, l_sl])
-        return out
+                out[m, k, l] = np_op(np_data[m_sl, k_sl, l_sl])
+    return out
 
-    if len(out_size) == 2:
-        pool_op = pool2d
-    else:
-        assert len(out_size) == 3
-        pool_op = pool3d
 
-    n, c = np_data.shape[:2]
+def adaptive_pool_nchw(np_data, out_size, pool_op, np_op):
+    """ The reference function for adaptive pool, nchw layout """
+    ishape = np_data.shape
+    n, c = ishape[:2]
     oshape = (n, c) + out_size
     np_out = np.zeros(oshape).astype(np_data.dtype)
-    np_op = np.mean if pool_type == "avg" else np.max
+
     for i in range(n):
         for j in range(c):
-            np_out[i, j] = pool_op(i, j)
+            np_out[i, j] = pool_op(ishape[2:], out_size, np_data[i, j], np_op)
+
     return np_out
+
+
+def adaptive_pool_nhwc(np_data, out_size, pool_op, np_op):
+    """ The reference function for adaptive pool, nhwc layout """
+    ishape = np_data.shape
+    n, c = ishape[0], ishape[-1]
+    oshape = (n,) + out_size + (c,)
+    np_out = np.zeros(oshape).astype(np_data.dtype)
+
+    for i in range(n):
+        for j in range(c):
+            if len(out_size) == 2:
+                np_out[i, :, :, j] = pool_op(ishape[1:-1], out_size,
+                                             np_data[i, :, :, j], np_op)
+            else:
+                np_out[i, :, :, :, j] = pool_op(ishape[1:-1], out_size,
+                                                np_data[i, :, :, :, j], np_op)
+
+    return np_out
+
+
+def adaptive_pool(np_data, out_size, pool_type, layout):
+    """ The reference function for adaptive pool, for 2d and 3d """
+    if len(out_size) == 2:
+        pool_op = _pool2d
+    else:
+        assert len(out_size) == 3
+        pool_op = _pool3d
+
+    np_op = np.mean if pool_type == "avg" else np.max
+
+    if layout in ["NCHW", "NCDHW"]:
+        return adaptive_pool_nchw(np_data, out_size, pool_op, np_op)
+
+    assert layout in ["NHWC", "NDHWC"]
+    return adaptive_pool_nhwc(np_data, out_size, pool_op, np_op)

topi/tests/python/test_topi_pooling.py

@@ -244,15 +244,10 @@ def test_global_pool():
     verify_global_pool(1, 1024, 7, 7, 'max', 'NHWC')
     verify_global_pool(4, 1024, 7, 7, 'max', 'NHWC')
 
-def verify_adaptive_pool(dshape, out_size, pool_type, layout="NCHW", dtype="float32"):
-    def start_index(index, odim, idim):
-        return int(np.floor(index * idim / odim))
-
-    def end_index(index, odim, idim):
-        return int(np.ceil((index + 1) * idim / odim))
 
+def verify_adaptive_pool(dshape, out_size, pool_type, layout="NCHW", dtype="float32"):
     np_data = np.random.uniform(low=0, high=255, size=dshape).astype(dtype)
-    np_out = topi.testing.adaptive_pool(np_data, out_size, pool_type)
+    np_out = topi.testing.adaptive_pool(np_data, out_size, pool_type, layout)
     oshape = np_out.shape
 
     data = te.placeholder(dshape, name="data", dtype=dtype)
@@ -280,15 +275,22 @@ def check_device(device):
     for device in get_all_backend():
         check_device(device)
 
+
 def test_adaptive_pool():
     verify_adaptive_pool((1, 3, 224, 224), (1, 1), "max")
     verify_adaptive_pool((1, 3, 224, 224), (1, 1), "avg")
     verify_adaptive_pool((1, 14, 56, 78), (34, 13), "max")
     verify_adaptive_pool((1, 5, 46, 97), (4, 96), "avg")
+    verify_adaptive_pool((1, 224, 224, 3), (1, 1), "max", layout="NHWC")
+    verify_adaptive_pool((1, 5, 46, 97), (4, 96), "avg", layout="NHWC")
     verify_adaptive_pool((1, 16, 32, 32, 32), (1, 1, 1), "max", layout="NCDHW")
     verify_adaptive_pool((1, 16, 32, 32, 32), (1, 1, 1), "avg", layout="NCDHW")
-    verify_adaptive_pool((1, 16, 32, 32, 32), (2, 2, 2), "max", layout="NCDHW")
     verify_adaptive_pool((1, 16, 32, 32, 32), (2, 2, 2), "avg", layout="NCDHW")
+    verify_adaptive_pool((1, 16, 64, 32, 32), (7, 8, 9), "avg", layout="NCDHW")
+    verify_adaptive_pool((1, 16, 64, 32, 32), (8, 16, 16), "avg", layout="NCDHW")
+    verify_adaptive_pool((1, 16, 32, 32, 32), (1, 1, 1), "avg", layout="NDHWC")
+    verify_adaptive_pool((1, 16, 32, 32, 32), (2, 2, 2), "max", layout="NDHWC")
+    verify_adaptive_pool((1, 16, 32, 32, 32), (2, 4, 4), "max", layout="NDHWC")
 
 
 def verify_pool3d(n, ic, ih, kh, sh, padding, pool_type,