Split DownsampleFactorMaxGrad

theano/sandbox/cuda/dnn.py

@@ -12,7 +12,7 @@ from theano.compile.ops import shape_i
 from theano.configparser import AddConfigVar, EnumStr
 from theano.tensor.nnet import SoftmaxGrad
 from theano.tensor.signal.downsample import (
-    DownsampleFactorMax, DownsampleFactorMaxGrad)
+    DownsampleFactorMax, MaxPoolGrad, AveragePoolGrad)
 from theano.sandbox.cuda import GpuOp
 from theano.sandbox.cuda.basic_ops import (as_cuda_ndarray_variable,
                                            host_from_gpu,
@@ -1672,11 +1672,11 @@ if True:
                                      desc)]
 
     @register_opt('cudnn')
-    @local_optimizer([DownsampleFactorMaxGrad])
+    @local_optimizer([MaxPoolGrad])
     def local_pool_dnn_grad_stride(node):
         if not dnn_available():
             return
-        if isinstance(node.op, DownsampleFactorMaxGrad):
+        if isinstance(node.op, MaxPoolGrad):
             if not node.op.ignore_border:
                 return
             inp, out, inp_grad = node.inputs
@@ -1696,6 +1696,31 @@ if True:
                                        desc)
                 return [host_from_gpu(ret)]
 
+    @register_opt('cudnn')
+    @local_optimizer([AveragePoolGrad])
+    def local_avgpool_dnn_grad_stride(node):
+        if not dnn_available():
+            return
+        if isinstance(node.op, AveragePoolGrad):
+            if not node.op.ignore_border:
+                return
+            inp, inp_grad = node.inputs
+            ds = node.op.ds
+            st = node.op.st
+            pad = node.op.padding
+            mode = node.op.mode
+
+            if ((inp.owner and isinstance(inp.owner.op, HostFromGpu)) or
+                (inp_grad.owner and isinstance(inp_grad.owner.op,
+                                               HostFromGpu))):
+                desc = GpuDnnPoolDesc(ws=ds, stride=st, mode=mode, pad=pad)()
+                ret = GpuDnnPoolGrad()(gpu_contiguous(inp),
+                                       gpu_contiguous(numpy.empty((1,1,1,1),
+                                           dtype=numpy.float32)),
+                                       gpu_contiguous(inp_grad),
+                                       desc)
+                return [host_from_gpu(ret)]
+
     @register_opt('cudnn')
     @local_optimizer([GpuSoftmax])
     def local_softmax_dnn(node):

theano/sandbox/cuda/opt.py

@@ -120,7 +120,8 @@ cpu_ops_moved_to_gpu = [
     tensor.blas.Dot22, tensor.blas.Dot22Scalar, tensor.blas.Gemm,
     tensor.blas.Gemv, tensor.blas.Ger, tensor.nnet.conv.ConvOp,
     tensor.signal.downsample.DownsampleFactorMax,
-    tensor.signal.downsample.DownsampleFactorMaxGrad,
+    tensor.signal.downsample.MaxPoolGrad,
+    tensor.signal.downsample.AveragePoolGrad,
     theano.tensor.nnet.neighbours.Images2Neibs,
     tensor.nnet.CrossentropySoftmaxArgmax1HotWithBias,
     tensor.nnet.CrossentropySoftmax1HotWithBiasDx,
@@ -1764,9 +1765,9 @@ def local_gpu_downsample_factor_max(node):
 
 
 @register_opt()
-@local_optimizer([downsample.DownsampleFactorMaxGrad])
+@local_optimizer([downsample.MaxPoolGrad])
 def local_gpu_downsample_factor_max_grad(node):
-    if (isinstance(node.op, downsample.DownsampleFactorMaxGrad) and
+    if (isinstance(node.op, downsample.MaxPoolGrad) and
         node.op.ds == node.op.st):
         assert node.op.__props__ == ('ds', 'ignore_border', 'st', 'padding',
                                      'mode')

theano/sandbox/cuda/tests/test_dnn.py

@@ -10,7 +10,7 @@ import theano.tensor as T
 import theano.tests.unittest_tools as utt
 from theano.sandbox.neighbours import images2neibs
 from theano.tensor.signal.downsample import max_pool_2d
-from theano.tensor.signal.downsample import DownsampleFactorMaxGrad
+from theano.tensor.signal.downsample import MaxPoolGrad, AveragePoolGrad
 import theano.sandbox.cuda.dnn as dnn
 from theano.sandbox.cuda.basic_ops import GpuAllocEmpty, gpu_alloc_empty
 
@@ -278,7 +278,11 @@ def test_pooling():
                               ignore_border=True, mode=mode)
             fc = theano.function([x], theano.grad(out.sum(), x),
                                  mode=mode_without_gpu)
-            assert any([isinstance(node.op, DownsampleFactorMaxGrad)
+            if mode == 'max':
+                assert any([isinstance(node.op, MaxPoolGrad)
+                            for node in fc.maker.fgraph.toposort()])
+            else:
+                assert any([isinstance(node.op, AveragePoolGrad)
                             for node in fc.maker.fgraph.toposort()])
             c_out = fc(data)
             assert numpy.allclose(c_out, g_out)

theano/tensor/signal/tests/test_downsample.py

@@ -8,7 +8,7 @@ import theano
 import theano.tensor as tensor
 from theano.tests import unittest_tools as utt
 from theano.tensor.signal.downsample import (DownsampleFactorMax, max_pool_2d,
-                                             DownsampleFactorMaxGrad,
+                                             MaxPoolGrad, AveragePoolGrad,
                                              DownsampleFactorMaxGradGrad,
                                              max_pool_2d_same_size)
 from theano import function
@@ -417,7 +417,7 @@ class TestDownsampleFactorMax(utt.InferShapeTester):
                 def mp(input, grad):
                     out = DownsampleFactorMax(
                         maxpoolshp, ignore_border=ignore_border)(input)
-                    grad_op = DownsampleFactorMaxGrad(
+                    grad_op = MaxPoolGrad(
                         maxpoolshp, ignore_border=ignore_border)
                     return grad_op(input, out, grad)
 
@@ -443,7 +443,7 @@ class TestDownsampleFactorMax(utt.InferShapeTester):
                         out = DownsampleFactorMax(
                             maxpoolshp, ignore_border=ignore_border,
                             st=stride)(input)
-                        grad_op = DownsampleFactorMaxGrad(
+                        grad_op = MaxPoolGrad(
                             maxpoolshp, ignore_border=ignore_border,
                             st=stride)
                         return grad_op(input, out, grad)
@@ -475,7 +475,7 @@ class TestDownsampleFactorMax(utt.InferShapeTester):
                     out = DownsampleFactorMax(
                         maxpoolshp, ignore_border=ignore_border,
                         st=stride)(input)
-                    grad_op = DownsampleFactorMaxGrad(
+                    grad_op = MaxPoolGrad(
                         maxpoolshp, ignore_border=ignore_border,
                         st=stride)
                     return grad_op(input, out, grad)
@@ -509,7 +509,7 @@ class TestDownsampleFactorMax(utt.InferShapeTester):
                     st=stridesize,
                     padding=paddingsize,
                     )(input)
-                grad_op = DownsampleFactorMaxGrad(maxpoolsize, ignore_border=True,
+                grad_op = MaxPoolGrad(maxpoolsize, ignore_border=True,
                                                   st=stridesize, padding=paddingsize)
                 return grad_op(input, out, grad)
             utt.verify_grad(mp, [imval, grad_val], rng=rng)
@@ -685,12 +685,12 @@ class TestDownsampleFactorMax(utt.InferShapeTester):
                     maxout_val = rng.rand(*out_shapes[k][i][j])
                     gz_val = rng.rand(*out_shapes[k][i][j])
                     self._compile_and_check([image, maxout, gz],
-                                            [DownsampleFactorMaxGrad(maxpoolshp,
+                                            [MaxPoolGrad(maxpoolshp,
                                             ignore_border=ignore_border,
                                             padding=padding)
                                             (image, maxout, gz)],
                                             [image_val, maxout_val, gz_val],
-                                            DownsampleFactorMaxGrad,
+                                            MaxPoolGrad,
                                             warn=False)