Move the softmax dnn test to test_dnn.py

theano/sandbox/cuda/tests/test_dnn.py

@@ -451,11 +451,98 @@ def test_pooling_opt():
                 for n in f.maker.fgraph.toposort()])
 
 
-def test_log_softmax():
+class test_DnnSoftMax(test_.test_SoftMax):
+    gpu_op = dnn.GpuDnnSoftmax
+    gpu_grad_op = dnn.GpuDnnSoftmaxGrad
+    mode = mode_with_gpu
+    topo_idx = -3
+
+    def setUp(self):
+        if not cuda.dnn.dnn_available():
+            raise SkipTest(cuda.dnn.dnn_available.msg)
+        utt.seed_rng()
+
+    def test_dnn_softmax_grad(self):
+        softmax_op = dnn.GpuDnnSoftmax('bc01', 'accurate', 'channel')
+
+        x_val = numpy.random.normal(0, 1, (3, 4, 2, 5)).astype('float32')
+        x_val2 = numpy.random.normal(0, 1, (3, 4, 1, 1)).astype('float32')
+
+        utt.verify_grad(softmax_op, [x_val])
+
+        utt.verify_grad(softmax_op, [x_val2])
+
+    def test_cudnn_softmax_grad_opt(self):
+        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad optimization is
+        # applied when cudnn is required
+        y = T.fvector('y')
+        f = theano.function(
+            [y],
+            T.grad(T.nnet.softmax(y).mean(), y),
+            mode=mode_with_gpu
+        )
+        sorted_f = f.maker.fgraph.toposort()
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
+                    )]) == 1)
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.tensor.nnet.SoftmaxGrad
+                    )]) == 0)
+
+        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad optimization is not
+        # applied when cudnn is excluded or not available
+        mode_wo_cudnn = mode_with_gpu.excluding("cudnn")
+        y = T.fvector('y')
+        f = theano.function(
+            [y],
+            T.grad(T.nnet.softmax(y).mean(), y),
+            mode=mode_wo_cudnn
+        )
+        sorted_f = f.maker.fgraph.toposort()
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
+                    )]) == 0)
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.tensor.nnet.SoftmaxGrad
+                    )]) == 1)
+
+        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad do not
+        # crash with manual graph
+        y = T.fvector('y')
+        o = theano.tensor.nnet.SoftmaxGrad()(y, y*2)
+        f = theano.function([y], o, mode=mode_with_gpu)
+        sorted_f = f.maker.fgraph.toposort()
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
+                    )]) == 1)
+        assert(len([i
+                    for i in sorted_f
+                    if isinstance(
+                        i.op,
+                        theano.tensor.nnet.SoftmaxGrad
+                    )]) == 0)
+
+
+    def test_log_softmax(self):
         # This is a test for an optimization that depends on CuDNN v3 or
         # more recent. Don't test if the CuDNN version is too old.
-    if not cuda.dnn.dnn_available() or cuda.dnn.version() < (3000, 3000):
-        raise SkipTest(cuda.dnn.dnn_available.msg)
+        if cuda.dnn.version() < (3000, 3000):
+            raise SkipTest("Log-softmax is only in cudnn v3+")
 
         x = T.ftensor4()
         softmax_out = dnn.GpuDnnSoftmax('bc01', 'accurate', 'channel')(x)
@@ -490,23 +577,6 @@ def test_log_softmax():
             utt.assert_allclose(out, expected_out)
 
 
-def test_dnn_softmax_grad():
-    utt.seed_rng()
-    softmax_op = dnn.GpuDnnSoftmax('bc01', 'accurate', 'channel')
-
-    x_val = numpy.random.normal(0, 1, (3, 4, 1, 1)).astype('float32')
-
-    utt.verify_grad(softmax_op, [x_val])
-
-
-def test_dnn_softmax_grad_opt():
-    utt.seed_rng()
-
-    x_val = numpy.random.normal(0, 1, (3, 4)).astype('float32')
-
-    utt.verify_grad(softmax_op, [x_val], mode=mode_with_gpu)
-
-
 def test_dnn_tag():
     """
     Test that if cudnn isn't avail we crash and that if it is avail, we use it.

theano/sandbox/cuda/tests/test_nnet.py

@@ -212,6 +212,12 @@ def test_softmax_with_bias():
 
 
 class test_SoftMax(unittest.TestCase):
+    gpu_op = cuda.nnet.GpuSoftmax
+    mode = mode_with_gpu.excluding("cudnn")
+    do_big = True
+    do_0 = True
+    topo_idx = -2
+
     def _test_softmax(
         self,
         x,
@@ -219,7 +225,6 @@ class test_SoftMax(unittest.TestCase):
         f_z,
         f_gpu_z,
         cmp,
-        gpu_mode,
         check_types
     ):
         """
@@ -232,7 +237,7 @@ class test_SoftMax(unittest.TestCase):
         f_gpu_z_out = f_gpu_z(x_gpu)
 
         f = theano.function([x], f_z_out, mode=mode_without_gpu)
-        f_gpu = theano.function([x_gpu], f_gpu_z_out, mode=gpu_mode)
+        f_gpu = theano.function([x_gpu], f_gpu_z_out, mode=self.mode)
         check_types(f, f_gpu)
 
         # we need to test n>32*1024 to check that we make the block loop.
@@ -261,16 +266,15 @@ class test_SoftMax(unittest.TestCase):
         return f, f_gpu
 
     def _cmp(self, n, m, f, f_gpu):
-        # print "test_softmax",n,m
         data = numpy.arange(n * m, dtype='float32').reshape(n, m)
         out = f(data)
         gout = f_gpu(data)
-        assert numpy.allclose(out, gout), numpy.absolute(out - gout)
+        utt.assert_allclose(out, gout)
 
-    def _check_types(self, graph, graph_gpu, topo_idx, f_type, f_gpu_type):
+    def _check_types(self, graph, graph_gpu, f_type, f_gpu_type):
         assert isinstance(graph.maker.fgraph.toposort()[-1].op, f_type)
         assert isinstance(
-            graph_gpu.maker.fgraph.toposort()[topo_idx].op,
+            graph_gpu.maker.fgraph.toposort()[self.topo_idx].op,
             f_gpu_type
         )
 
@@ -278,180 +282,24 @@ class test_SoftMax(unittest.TestCase):
         x = T.fmatrix('x')
         z = T.nnet.softmax_op
 
-        def check_types_without_cudnn(graph, graph_gpu):
+        def check_types(graph, graph_gpu):
             self._check_types(
                 graph,
                 graph_gpu,
-                -2,
                 type(z),
-                cuda.nnet.GpuSoftmax
+                self.gpu_op
             )
 
-        mode_wo_cudnn = mode_with_gpu.excluding("cudnn")
         f, f_gpu = self._test_softmax(
             x,
             x,
             z,
             z,
             self._cmp,
-            mode_wo_cudnn,
-            check_types_without_cudnn
+            check_types
         )
 
-        # cuDNN R1 cannot handle these test cases but the Theano softmax can so
-        # we test them only for the Theano softmax.
+        if do_big:
             self._cmp(2 << 15, 5, f, f_gpu)
+        if do_0:
             self._cmp(0, 10, f, f_gpu)
-
-    def test_softmax_cudnn(self):
-        if not cuda.dnn.dnn_available():
-            raise SkipTest(cuda.dnn.dnn_available.msg)
-        x = T.fmatrix('x')
-        z = T.nnet.softmax_op
-
-        def check_types_with_cudnn(graph, graph_gpu):
-            self._check_types(
-                graph,
-                graph_gpu,
-                -3,
-                type(z),
-                theano.sandbox.cuda.dnn.GpuDnnSoftmax
-            )
-
-        f, f_gpu = self._test_softmax(
-            x,
-            x,
-            z,
-            z,
-            self._cmp,
-            mode_with_gpu,
-            check_types_with_cudnn
-        )
-
-    def test_cudnn_softmax_grad(self):
-        if not cuda.dnn.dnn_available():
-            raise SkipTest(cuda.dnn.dnn_available.msg)
-
-        def cmp(n, m, f, f_gpu):
-            data = numpy.arange(n * m, dtype='float32').reshape(n, m)
-            gdata = numpy.asarray(data)[:, :, None, None]
-            out = f(data)
-            gout = numpy.asarray(f_gpu(gdata))[:, :, 0, 0]
-            assert numpy.allclose(out, gout), numpy.absolute(out - gout)
-
-        x = T.matrix('x', 'float32')
-        x_gpu = T.tensor4('x_gpu', 'float32')
-        f_z = T.nnet.softmax_op
-        f_gpu = theano.sandbox.cuda.dnn.GpuDnnSoftmax(
-            'bc01',
-            'accurate',
-            'channel'
-        )
-
-        # Verify the grad operation
-        dims = (2, 3, 4, 5)
-        gdata = numpy.arange(
-            numpy.product(dims),
-            dtype='float32'
-        ).reshape(dims)
-        T.verify_grad(f_gpu, [gdata], rng=numpy.random,
-                      mode=mode_with_gpu)
-
-        def check_types(graph, graph_gpu):
-            self._check_types(
-                graph,
-                graph_gpu,
-                -1,
-                type(f_z),
-                theano.sandbox.cuda.dnn.GpuDnnSoftmax
-            )
-
-        def check_types_opt(graph, graph_gpu):
-            assert isinstance(graph.maker.fgraph.toposort()[-1].op, type(f_z))
-            assert len([n for n in graph_gpu.maker.fgraph.toposort()
-                        if isinstance(
-                            n.op,
-                            theano.sandbox.cuda.dnn.GpuDnnSoftmax
-                        )]) == 1
-
-        # Verify that the CPU and GPU implementations return the same results
-        # up to a tolerance.
-        self._test_softmax(
-            x,
-            x_gpu,
-            f_z,
-            f_gpu,
-            cmp,
-            mode_with_gpu,
-            check_types
-        )
-
-        mode_w_cudnn = mode_with_gpu.including("cudnn")
-        self._test_softmax(
-            x, x, f_z, f_z, self._cmp,
-            mode_w_cudnn, check_types_opt
-        )
-
-        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad optimization is
-        # applied when cudnn is required
-        y = T.fvector('y')
-        f = theano.function(
-            [y],
-            T.grad(T.nnet.softmax(y).mean(), y),
-            mode=mode_with_gpu
-        )
-        sorted_f = f.maker.fgraph.toposort()
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
-                    )]) == 1)
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.tensor.nnet.SoftmaxGrad
-                    )]) == 0)
-
-        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad optimization is not
-        # applied when cudnn is excluded or not available
-        mode_wo_cudnn = mode_with_gpu.excluding("cudnn")
-        y = T.fvector('y')
-        f = theano.function(
-            [y],
-            T.grad(T.nnet.softmax(y).mean(), y),
-            mode=mode_wo_cudnn
-        )
-        sorted_f = f.maker.fgraph.toposort()
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
-                    )]) == 0)
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.tensor.nnet.SoftmaxGrad
-                    )]) == 1)
-
-        # Verify that the SoftmaxGrad -> GpuDnnSoftmaxGrad do not
-        # crash with manual graph
-        y = T.fvector('y')
-        o = theano.tensor.nnet.SoftmaxGrad()(y, y*2)
-        f = theano.function([y], o, mode=mode_with_gpu)
-        sorted_f = f.maker.fgraph.toposort()
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.sandbox.cuda.dnn.GpuDnnSoftmaxGrad
-                    )]) == 1)
-        assert(len([i
-                    for i in sorted_f
-                    if isinstance(
-                        i.op,
-                        theano.tensor.nnet.SoftmaxGrad
-                    )]) == 0)