fix test in float32, but also check float64 at the same time.

theano/sandbox/cuda/tests/test_memory.py

@@ -18,7 +18,7 @@ else:
     mode_with_gpu = theano.compile.mode.get_default_mode().including('gpu')
 
 
-def freemem():
+def freemem(extra_alloc=0):
     """
     Return the free memory on the gpu in megabytes.
     """
@@ -30,15 +30,16 @@ def freemem():
     if hasattr(cuda.cuda_ndarray.cuda_ndarray, "theano_allocated"):
         theano_alloc = cuda.cuda_ndarray.cuda_ndarray.theano_allocated()
         return ("(n malloc/theano mem allocated in KB)",
-                n_mallocs, int(theano_alloc / 1024))
+                n_mallocs + extra_alloc,
+                int(theano_alloc / 1024) + extra_size)
 
-    return ("n malloc on the gpu", n_mallocs)
+    return ("n malloc on the gpu", n_mallocs + extra_alloc)
     # I don't use the following by default as if there is other stuff running
     # on the GPU, this won't work.
     mem_info = cuda.cuda_ndarray.cuda_ndarray.mem_info()
     gpu_used = (mem_info[1] - mem_info[0]) / 1024 ** 2
     mem_info_msg = "(n malloc/gpu mem used in MB)"
-    return ("(n malloc/gpu mem used in MB)", n_mallocs, int(gpu_used))
+    return (mem_info_msg, n_mallocs, int(gpu_used))
 
 
 def test_memory():
@@ -52,15 +53,26 @@ def test_memory():
     note::
         This test can fail if there is other process running on the gpu.
     """
-    shapes = (6000, 5000)
-    test_params = np.asarray(np.random.randn(np.prod(shapes)), 'float32')
-
-    some_vector = tensor.vector('some_vector')
+    shapes = (200, 100)
+    # more_alloc1 and more_alloc2 is not the same for both dtype.
+    # when dtype is float32, the computation is done on the gpu.
+    # This insert constant on the gpu during compilation
+    # that raise the number of alloc.
+    # When dtype is float64, only the shared is on the gpu and it is transferd
+    # to the cpu for computation. So no extra alloc after compilation.
+    # more_alloc1 if after the first compilation, more_alloc2 after the second.
+    for dtype, more_alloc1, more_alloc2 in [("float32", 2, 9),
+                                            ("float64", 0, 0)]:
+        print dtype
+        test_params = np.asarray(np.random.randn(np.prod(shapes)), dtype)
+
+        some_vector = tensor.vector('some_vector', dtype=dtype)
         some_matrix = some_vector.reshape(shapes)
 
         mem1 = freemem()
         print "Before shared variable", mem1
-    variables = cuda.shared_constructor(np.ones((shapes[1],), dtype='float32'))
+        variables = cuda.shared_constructor(np.ones((shapes[1],),
+                                                    dtype='float32'))
         derp = tensor.sum(tensor.dot(some_matrix[:shapes[0]], variables))
         print "Shared took ", np.prod(variables.get_value(
                 borrow=True,
@@ -68,28 +80,30 @@ def test_memory():
 
         mem2 = freemem()
         print "Before compilation", mem2
+        mem2_1 = freemem(extra_alloc=more_alloc1)
+        mem2_2 = freemem(extra_alloc=more_alloc2)
         obj = theano.function([some_vector], derp, mode=mode_with_gpu)
         mem3 = freemem()
         print "After function compilation 1", mem3
-    assert mem2 == mem3, (mem2, mem3)
+        assert mem2_1 == mem3, (mem2_1, mem3)
 
         grad_derp = tensor.grad(derp, some_vector)
         grad = theano.function([some_vector], grad_derp, mode=mode_with_gpu)
         mem4 = freemem()
         print "After function compilation 2", mem4
-    assert mem2 == mem4, (mem2, mem4)
+        assert mem2_2 == mem4, (mem2_2, mem4)
 
         for i in range(3):
             obj(test_params)
             print "After function evaluation 1", freemem()
-        assert mem2 == freemem(), (mem2, freemem())
+            assert mem2_2 == freemem(), (mem2_2, freemem())
             grad(test_params)
             print "After function evaluation 2", freemem()
-        assert mem2 == freemem(), (mem2, freemem())
+            assert mem2_2 == freemem(), (mem2_2, freemem())
 
         del obj
-    print "After deleting function 1", freemem()
-    assert mem2 == freemem(), (mem2, freemem())
+        #print "After deleting function 1", freemem()
+        #assert mem2 == freemem(), (mem2, freemem())
 
         del grad
         print "After deleting function 2", freemem()