Merge pull request #732 from larseeri/shape_downsample

theano/tensor/signal/downsample.py

@@ -65,7 +65,8 @@ class DownsampleFactorMax(Op):
 
         :param imgshape: the shape of a tensor of images. The last two elements are interpreted
         as the number of rows, and the number of cols.
-        :type imgshape: tuple, list, or similar.
+        :type imgshape: tuple, list, or similar of integer or
+        scalar Theano variable.
 
         :param ds: downsample factor over rows and columns
         :type ds: list or tuple of two ints
@@ -83,10 +84,15 @@ class DownsampleFactorMax(Op):
             raise TypeError('imgshape must have at least two elements (rows, cols)')
         r, c = imgshape[-2:]
         rval = list(imgshape[:-2])+[ r/ds[0], c/ds[1]]
+
         if not ignore_border:
-            if r % ds[0]:
+            if isinstance(r, theano.Variable):
+                rval[-2] = tensor.switch(r % ds[0], rval[-2] + 1, rval[-2])
+            elif r % ds[0]:
                 rval[-2] += 1
-            if c % ds[1]:
+            if isinstance(c, theano.Variable):
+                rval[-1] = tensor.switch(c % ds[1], rval[-1] + 1, rval[-1])
+            elif c % ds[1]:
                 rval[-1] += 1
         return rval
 
@@ -149,6 +155,10 @@ class DownsampleFactorMax(Op):
                         zj = j / ds1
                         zz[n,k,zi,zj] = __builtin__.max(zz[n,k,zi,zj], x[n,k,i,j])
 
+    def infer_shape(self, node, in_shapes):
+        shp = self.out_shape(in_shapes[0], self.ds, self.ignore_border)
+        return [shp]
+
     def grad(self, inp, grads):
         x, = inp
         gz, = grads
@@ -275,6 +285,9 @@ class DownsampleFactorMaxGrad(Op):
                         else: gx[n,k,i,j] = 0
         gx_stg[0] = gx
 
+    def infer_shape(self, node, in_shapes):
+        return [in_shapes[0]]
+
     def c_code(self, node, name, inp, out, sub):
         x, z, gz = inp
         gx, = out

theano/tensor/signal/tests/test_downsample.py

-import unittest, sys, time
+import unittest
+import sys
+import time
 import numpy
 import theano.tensor as tensor
 from theano.tests import unittest_tools as utt
-from theano.tensor.signal.downsample import DownsampleFactorMax, max_pool_2d
+from theano.tensor.signal.downsample import (DownsampleFactorMax, max_pool_2d,
+                                             DownsampleFactorMaxGrad)
 from theano import function, Mode
 
 
-class TestDownsampleFactorMax(unittest.TestCase):
-    def setUp(self):
-        utt.seed_rng()
+class TestDownsampleFactorMax(utt.InferShapeTester):
 
     @staticmethod
     def numpy_max_pool_2d(input, ds, ignore_border=False):
         '''Helper function, implementing max_pool_2d in pure numpy'''
         if len(input.shape) < 2:
-            raise NotImplementedError('input should have at least 2 dim, shape is %s'\
+            raise NotImplementedError('input should have at least 2 dim,'
+                                      ' shape is %s'\
                     % str(input.shape))
-
-        xi=0
-        yi=0
+        xi = 0
+        yi = 0
         if not ignore_border:
             if input.shape[-2] % ds[0]:
                 xi += 1
             if input.shape[-1] % ds[1]:
                 yi += 1
-
         out_shp = list(input.shape[:-2])
-        out_shp.append(input.shape[-2]/ds[0]+xi)
-        out_shp.append(input.shape[-1]/ds[1]+yi)
-
+        out_shp.append(input.shape[-2] / ds[0] + xi)
+        out_shp.append(input.shape[-1] / ds[1] + yi)
         output_val = numpy.zeros(out_shp)
 
         for k in numpy.ndindex(input.shape[:-2]):
             for i in range(output_val.shape[-2]):
-                ii =  i*ds[0]
+                ii = i * ds[0]
                 for j in range(output_val.shape[-1]):
-                    jj = j*ds[1]
-                    patch = input[k][ii:ii+ds[0],jj:jj+ds[1]]
-                    output_val[k][i,j] = numpy.max(patch)
+                    jj = j * ds[1]
+                    patch = input[k][ii:ii + ds[0], jj:jj + ds[1]]
+                    output_val[k][i, j] = numpy.max(patch)
         return output_val
 
     def test_DownsampleFactorMax(self):
         rng = numpy.random.RandomState(utt.fetch_seed())
-
         # generate random images
-        maxpoolshps = ((1,1),(2,2),(3,3),(2,3))
-        imval = rng.rand(4,10,64,64)
+        maxpoolshps = ((1, 1), (2, 2), (3, 3), (2, 3))
+        imval = rng.rand(4, 10, 64, 64)
         images = tensor.dtensor4()
 
         for maxpoolshp in maxpoolshps:
-            for ignore_border in [True,False]:
+            for ignore_border in [True, False]:
                 #print 'maxpoolshp =', maxpoolshp
                 #print 'ignore_border =', ignore_border
 
-                ## Pure Numpy computation
-                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp, ignore_border)
-
+                # Pure Numpy computation
+                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp,
+                                                          ignore_border)
                 output = max_pool_2d(images, maxpoolshp, ignore_border)
-                f = function([images,],[output,])
+                f = function([images, ], [output, ])
                 output_val = f(imval)
                 assert numpy.all(output_val == numpy_output_val)
 
                 #DownsampleFactorMax op
-                maxpool_op = DownsampleFactorMax(maxpoolshp, ignore_border=ignore_border)(images)
+                maxpool_op = DownsampleFactorMax(maxpoolshp,
+                                                 ignore_border=ignore_border)(images)
                 f = function([images], maxpool_op)
                 output_val = f(imval)
                 assert (numpy.abs(output_val - numpy_output_val) < 1e-5).all()
 
     def test_DownsampleFactorMax_grad(self):
         rng = numpy.random.RandomState(utt.fetch_seed())
-        maxpoolshps = ((1,1),(3,2),(2,3))
-        imval = rng.rand(2,3,3,4) * 10.0 #more variance means numeric gradient will be more accurate
+        maxpoolshps = ((1, 1), (3, 2), (2, 3))
+        imval = rng.rand(2, 3, 3, 4) * 10.0
+        #more variance means numeric gradient will be more accurate
 
         for maxpoolshp in maxpoolshps:
-            for ignore_border in [True,False]:
+            for ignore_border in [True, False]:
                 #print 'maxpoolshp =', maxpoolshp
                 #print 'ignore_border =', ignore_border
                 def mp(input):
-                    return DownsampleFactorMax(maxpoolshp, ignore_border=ignore_border)(input)
+                    return DownsampleFactorMax(maxpoolshp,
+                                    ignore_border=ignore_border)(input)
                 utt.verify_grad(mp, [imval], rng=rng)
 
     def test_max_pool_2d_2D(self):
         rng = numpy.random.RandomState(utt.fetch_seed())
-
-        maxpoolshps = ((1,1),(3,2))
-        imval = rng.rand(4,5)
+        maxpoolshps = ((1, 1), (3, 2))
+        imval = rng.rand(4, 5)
         images = tensor.dmatrix()
 
         for maxpoolshp in maxpoolshps:
-            for ignore_border in [True,False]:
+            for ignore_border in [True, False]:
                 #print 'maxpoolshp =', maxpoolshp
                 #print 'ignore_border =', ignore_border
-                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp, ignore_border)
-
+                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp,
+                                                          ignore_border)
                 output = max_pool_2d(images, maxpoolshp, ignore_border)
                 output_val = function([images], output)(imval)
                 assert numpy.all(output_val == numpy_output_val)
-
                 def mp(input):
                     return max_pool_2d(input, maxpoolshp, ignore_border)
                 utt.verify_grad(mp, [imval], rng=rng)
 
     def test_max_pool_2d_3D(self):
         rng = numpy.random.RandomState(utt.fetch_seed())
-
-        maxpoolshps = [(1,2)]
-        imval = rng.rand(2,3,4)
+        maxpoolshps = [(1, 2)]
+        imval = rng.rand(2, 3, 4)
         images = tensor.dtensor3()
 
         for maxpoolshp in maxpoolshps:
-            for ignore_border in [True,False]:
+            for ignore_border in [True, False]:
                 #print 'maxpoolshp =', maxpoolshp
                 #print 'ignore_border =', ignore_border
-                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp, ignore_border)
-
+                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp,
+                                                          ignore_border)
                 output = max_pool_2d(images, maxpoolshp, ignore_border)
                 output_val = function([images], output)(imval)
                 assert numpy.all(output_val == numpy_output_val)
-
                 c = tensor.sum(output)
                 c_val = function([images], c)(imval)
-
                 g = tensor.grad(c, images)
                 g_val = function([images],
                         [g.shape,
-                            tensor.min(g, axis=(0,1,2)),
-                            tensor.max(g, axis=(0,1,2))]
+                            tensor.min(g, axis=(0, 1, 2)),
+                            tensor.max(g, axis=(0, 1, 2))]
                         )(imval)
 
 #removed as already tested in test_max_pool_2d_2D
@@ -134,20 +130,18 @@ class TestDownsampleFactorMax(unittest.TestCase):
 #                    return max_pool_2d(input, maxpoolshp, ignore_border)
 #                utt.verify_grad(mp, [imval], rng=rng)
 
-
     def test_max_pool_2d_6D(self):
         rng = numpy.random.RandomState(utt.fetch_seed())
-
-        maxpoolshps = [(3,2)]
-        imval = rng.rand(2,1,1,1,3,4)
-        images = tensor.TensorType('float64', [False]*6)()
+        maxpoolshps = [(3, 2)]
+        imval = rng.rand(2, 1, 1, 1, 3, 4)
+        images = tensor.TensorType('float64', [False] * 6)()
 
         for maxpoolshp in maxpoolshps:
-            for ignore_border in [True,False]:
+            for ignore_border in [True, False]:
                 #print 'maxpoolshp =', maxpoolshp
                 #print 'ignore_border =', ignore_border
-                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp, ignore_border)
-
+                numpy_output_val = self.numpy_max_pool_2d(imval, maxpoolshp,
+                                                          ignore_border)
                 output = max_pool_2d(images, maxpoolshp, ignore_border)
                 output_val = function([images], output)(imval)
                 assert numpy.all(output_val == numpy_output_val)
@@ -158,6 +152,39 @@ class TestDownsampleFactorMax(unittest.TestCase):
 #                    return max_pool_2d(input, maxpoolshp, ignore_border)
 #                utt.verify_grad(mp, [imval], rng=rng)
 
+    def test_infer_shape(self):
+        image = tensor.dtensor4()
+        maxout = tensor.dtensor4()
+        gz = tensor.dtensor4()
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        maxpoolshps = ((1, 1), (2, 2), (3, 3), (2, 3), (3, 2))
+
+        image_val = rng.rand(4, 6, 7, 9)
+        out_shapes = [[[4, 6, 7, 9], [4, 6, 7, 9]],
+                      [[4, 6, 3, 4], [4, 6, 4, 5]],
+                      [[4, 6, 2, 3], [4, 6, 3, 3]],
+                      [[4, 6, 3, 3], [4, 6, 4, 3]],
+                      [[4, 6, 2, 4], [4, 6, 3, 5]]]
+
+        for i, maxpoolshp in enumerate(maxpoolshps):
+            for j, ignore_border in enumerate([True, False]):
+
+                # checking shapes generated by DownsampleFactorMax
+                self._compile_and_check([image],
+                        [DownsampleFactorMax(maxpoolshp,
+                        ignore_border=ignore_border)(image)],
+                        [image_val], DownsampleFactorMax)
+
+                # checking shapes generated by DownsampleFactorMaxGrad
+                maxout_val = rng.rand(*out_shapes[i][j])
+                gz_val = rng.rand(*out_shapes[i][j])
+                self._compile_and_check([image, maxout, gz],
+                        [DownsampleFactorMaxGrad(maxpoolshp,
+                        ignore_border=ignore_border)(image, maxout, gz)],
+                        [image_val, maxout_val, gz_val],
+                                        DownsampleFactorMaxGrad)
+
+
 
 
 if __name__ == '__main__':