Add tests for pooling and pooling grad

theano/gpuarray/tests/test_blas.py

@@ -12,7 +12,8 @@ from theano import tensor
 from theano.tests import unittest_tools as utt
 from theano.tensor.blas import gemv_inplace, gemm_inplace, _dot22, batched_dot
 from theano.tensor.tests.test_blas import TestGer, BaseGemv
-from theano.tensor.signal.pool import Pool, DownsampleFactorMaxGradGrad
+from theano.tensor.signal.pool import (Pool, MaxPoolGrad, AveragePoolGrad,
+                                       DownsampleFactorMaxGradGrad)
 
 from .. import gpuarray_shared_constructor
 from .config import mode_with_gpu, mode_without_gpu
@@ -21,7 +22,8 @@ from .test_basic_ops import makeTester, rand
 from ..blas import (gpugemv_inplace, gpugemv_no_inplace,
                     gpugemm_inplace, gpugemmbatch_no_inplace,
                     gpuger_inplace, gpuger_no_inplace,
-                    GpuGer, gpu_dot22, GpuDownsampleFactorMaxGradGrad)
+                    GpuGer, gpu_dot22, GpuPool, GpuMaxPoolGrad,
+                    GpuAveragePoolGrad, GpuDownsampleFactorMaxGradGrad)
 
 
 GpuGemvTester = makeTester(
@@ -135,7 +137,7 @@ GpuDot22Tester = makeTester(
 )
 
 
-def test_max_pool2d_grad_grad():
+def test_pool2d():
     shps = [(1, 12),
             (1, 1, 12),
             (1, 1, 1, 12),
@@ -166,28 +168,64 @@ def test_max_pool2d_grad_grad():
             (1, 1, 1023, 10), ]
 
     numpy.random.RandomState(utt.fetch_seed()).shuffle(shps)
-    test_ds = (2, 2), (3, 2), (1, 1)
+    test_ws = (2, 2), (3, 2), (1, 1)
     test_st = (2, 2), (3, 2), (1, 1)
+    test_mode = ['max', 'sum', 'average_inc_pad', 'average_exc_pad']
+
+    ref_mode = copy.copy(mode_without_gpu)
+    ref_mode.check_py_code = False
+    gpu_mode = copy.copy(mode_with_gpu)
+    gpu_mode.check_py_code = False
 
     for shp in shps:
-        for ds, st in itertools.product(test_ds, test_st):
-            if ds[0] > shp[-2] or ds[1] > shp[-1]:
+        for mode, ws, st in itertools.product(test_mode, test_ws, test_st):
+            if ws[0] > shp[-2] or ws[1] > shp[-1]:
                 continue
             for ignore_border, pad in zip((True, False), [(1, 1), (0, 0)]):
-                if pad[0] >= ds[0] or pad[1] >= ds[1]:
+                if pad[0] >= ws[0] or pad[1] >= ws[1]:
                     continue
-                # print('test_downsample', shp, ds, st, pad, ignore_border)
-                ds_op = Pool(ndim=len(ds), ignore_border=ignore_border)
+                if mode == 'average_exc_pad' and (pad[0] > 0 or pad[1] > 0):
+                    continue
+                print('test_pool2d', shp, ws, st, pad, mode, ignore_border)
+                ds_op = Pool(ndim=len(ws), mode=mode, ignore_border=ignore_border)
 
                 a = theano.shared(rand(*shp), 'a')
+                a_pooled = ds_op(tensor.as_tensor_variable(a), ws, st, pad)
+
+                f = theano.function([], a_pooled, mode=gpu_mode)
+                f2 = theano.function([], a_pooled, mode=ref_mode)
+
+                assert any([isinstance(node.op, GpuPool)
+                            for node in f.maker.fgraph.toposort()])
+                assert any([isinstance(node.op, Pool)
+                            for node in f2.maker.fgraph.toposort()])
+                assert numpy.allclose(f(), f2()), (shp, ws, st, pad, mode, ignore_border)
+
+                a_pooled_grad = tensor.grad(a_pooled.sum(), a)
+
+                g = theano.function([], a_pooled_grad, mode=gpu_mode)
+                g2 = theano.function([], a_pooled_grad, mode=ref_mode)
+
+                if mode == 'max':
+                    gop = GpuMaxPoolGrad
+                    gop2 = MaxPoolGrad
+                else:
+                    gop = GpuAveragePoolGrad
+                    gop2 = AveragePoolGrad
+                assert any([isinstance(node.op, gop)
+                            for node in g.maker.fgraph.toposort()])
+                assert any([isinstance(node.op, gop2)
+                            for node in g2.maker.fgraph.toposort()])
+
+                assert numpy.allclose(g(), g2()), (shp, ws, st, pad, mode, ignore_border)
+
+                # test grad grad for max pooling
+                # for average pooling grad grad is just average pooling grad
+                if mode != 'max':
+                    continue
 
-                ggf = gradient.Lop(tensor.grad((ds_op(
-                    tensor.as_tensor_variable(a), ds, st, pad)**2).sum(), a), a, a)
+                ggf = gradient.Lop(tensor.grad((a_pooled**2).sum(), a), a, a)
 
-                ref_mode = copy.copy(mode_without_gpu)
-                ref_mode.check_py_code = False
-                gpu_mode = copy.copy(mode_with_gpu)
-                gpu_mode.check_py_code = False
                 gg = theano.function([], ggf, mode=gpu_mode)
                 gg2 = theano.function([], ggf, mode=ref_mode)
 
@@ -199,7 +237,7 @@ def test_max_pool2d_grad_grad():
                     isinstance(node.op, DownsampleFactorMaxGradGrad)
                     for node in gg2.maker.fgraph.toposort()
                 ])
-                assert numpy.allclose(gg(), gg2()), (shp, ds, st,
+                assert numpy.allclose(gg(), gg2()), (shp, ws, st,
                                                      ignore_border)