Repository: pytorch/pytorch

Issue:
test_cudnn_weight_format nondeterministically fails when run with pytest
To reproduce, install pytest and run: `pytest test/test_nn.py -k cudnn`

You will get:

```
test/test_nn.py:2483: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
test/common.py:245: in assertEqual
    super(TestCase, self).assertLessEqual(abs(x - y), prec, message)
E   AssertionError: 1 not less than or equal to 1e-05 :

 test/test_nn.py ⨯✓                                                         100% ██████████
==================================== warnings summary =====================================
test/test_nn.py
  cannot collect 'test' because it is not a function.

test/test_nn.py::TestBase
  cannot collect test class 'TestBase' because it has a __init__ constructor

-- Docs: http://doc.pytest.org/en/latest/warnings.html
================================== 655 tests deselected ===================================

Results (16.94s):
       7 passed
       1 failed
         - test/test_nn.py:2450 TestNN.test_cudnn_weight_format
     655 deselected
```

This failure does not happen when you run only this test.


Discussion:
Minimal fix:

```
diff --git a/test/test_nn.py b/test/test_nn.py
index 46acdca..0cbd6a4 100644
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -1951,24 +1951,25 @@ class TestNN(NNTestCase):
     def test_Conv2d_inconsistent_types_on_GPU_without_cudnn(self):
         inputs = Variable(torch.randn(4, 1, 7, 7).float().cuda())
         weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
         bias = Variable(torch.randn(1).double().cuda())
 
         torch.backends.cudnn.enabled = False
         # inconsistent types should raise an exception
         self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
         self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
 
         # but it should work with the same type
         nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
+        torch.backends.cudnn.enabled = True
```

More robust fix coming tomorrow.

Resolving patch:
diff --git a/torch/backends/cudnn/__init__.py b/torch/backends/cudnn/__init__.py
--- a/torch/backends/cudnn/__init__.py
+++ b/torch/backends/cudnn/__init__.py
@@ -18,6 +18,13 @@
 # TODO: dynamic version checks via cudnnGetVersion
 
 
+# The idea for this parameter is that we forbid bare assignment
+# to torch.backends.cudnn.enabled and friends when running our
+# test suite, where it's very easy to forget to undo the change
+# later.
+__allow_nonbracketed_mutation_flag = True
+
+
 def find_cudnn_windows_lib():
     proc = Popen(['where', 'cudnn64*.dll'], stdout=PIPE, stderr=PIPE)
     out, err = proc.communicate()
@@ -122,14 +129,36 @@ def set_flags(_enabled, _benchmark, _deterministic, _verbose):
     return orig_flags
 
 
+def disable_global_flags():
+    global __allow_nonbracketed_mutation_flag
+    __allow_nonbracketed_mutation_flag = False
+
+
+def flags_frozen():
+    return not __allow_nonbracketed_mutation_flag
+
+
+@contextmanager
+def __allow_nonbracketed_mutation():
+    global __allow_nonbracketed_mutation_flag
+    old = __allow_nonbracketed_mutation_flag
+    __allow_nonbracketed_mutation_flag = True
+    try:
+        yield
+    finally:
+        __allow_nonbracketed_mutation_flag = old
+
+
 @contextmanager
 def flags(enabled=False, benchmark=False, deterministic=False, verbose=False):
-    orig_flags = set_flags(enabled, benchmark, deterministic, verbose)
+    with __allow_nonbracketed_mutation():
+        orig_flags = set_flags(enabled, benchmark, deterministic, verbose)
     try:
         yield
     finally:
         # recover the previous values
-        set_flags(orig_flags[0], orig_flags[1], orig_flags[2], orig_flags[3])
+        with __allow_nonbracketed_mutation():
+            set_flags(orig_flags[0], orig_flags[1], orig_flags[2], orig_flags[3])
 
 
 class CuDNNHandle:
@@ -397,7 +426,11 @@ def __get__(self, obj, objtype):
         return self.getter()
 
     def __set__(self, obj, val):
-        self.setter(val)
+        if not flags_frozen():
+            self.setter(val)
+        else:
+            raise RuntimeError("not allowed to set torch.backends.cudnn flags "
+                               "after disable_global_flags; please use flags() context manager instead")
 
 
 class CudnnModule(object):


Test patch:
diff --git a/test/common.py b/test/common.py
--- a/test/common.py
+++ b/test/common.py
@@ -17,9 +17,11 @@
 import torch.cuda
 from torch.autograd import Variable
 from torch._six import string_classes
+import torch.backends.cudnn
 
 
 torch.set_default_tensor_type('torch.DoubleTensor')
+torch.backends.cudnn.disable_global_flags()
 
 
 parser = argparse.ArgumentParser(add_help=False)
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -200,14 +200,11 @@ def _do_test(self, test_case, module, input):
 
                 # test that forwards of module runs correctly without cuDNN
                 if self.cudnn:
-                    torch.backends.cudnn.enabled = False
-                    try:
+                    with torch.backends.cudnn.flags(enabled=False):
                         module(input)
                         for p in module.parameters():
                             test_case.assertEqual(type(p.data), torch.cuda.FloatTensor)
                             test_case.assertEqual(p.get_device(), 0)
-                    finally:
-                        torch.backends.cudnn.enabled = True
 
                 if torch.cuda.device_count() >= 2:
                     # test cross-GPU transfer works
@@ -2037,13 +2034,13 @@ def test_Conv2d_inconsistent_types_on_GPU_without_cudnn(self):
         weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
         bias = Variable(torch.randn(1).double().cuda())
 
-        torch.backends.cudnn.enabled = False
-        # inconsistent types should raise an exception
-        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
-        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
+        with torch.backends.cudnn.flags(enabled=False):
+            # inconsistent types should raise an exception
+            self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
+            self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
 
-        # but it should work with the same type
-        nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
+            # but it should work with the same type
+            nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
 
     @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
     @unittest.skipIf(not TEST_CUDNN, 'CUDNN not available')
@@ -2052,13 +2049,13 @@ def test_Conv2d_inconsistent_types_on_GPU_with_cudnn(self):
         weights = Variable(torch.randn(1, 1, 3, 3).double().cuda())
         bias = Variable(torch.randn(1).double().cuda())
 
-        torch.backends.cudnn.enabled = True
-        # inconsistent types should raise an exception
-        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
-        self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
+        with torch.backends.cudnn.flags(enabled=True):
+            # inconsistent types should raise an exception
+            self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights))
+            self.assertRaises(RuntimeError, lambda: nn.functional.conv2d(inputs, weights.float(), bias))
 
-        # but it should work with the same type
-        nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
+            # but it should work with the same type
+            nn.functional.conv2d(inputs.float(), weights.float(), bias.float())
 
     @unittest.skipIf(IS_WINDOWS, 'TODO: fix this test for Windows')
     @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
@@ -2104,19 +2101,15 @@ def test_Conv2d_large_workspace(self):
         dtype = torch.cuda.FloatTensor
 
         def run_test(benchmark):
-            torch.backends.cudnn.benchmark = benchmark
-            conv = torch.nn.Conv2d(256, 256, kernel_size=3, padding=1).type(dtype)
-            for size in sizes:
-                x = torch.randn(size).type(dtype)
-                out = conv(Variable(x, requires_grad=True))
-                out.backward(torch.ones(out.size()).type(dtype))
-
-        b = torch.backends.cudnn.benchmark
-        try:
-            run_test(benchmark=False)
-            run_test(benchmark=True)
-        finally:
-            torch.backends.cudnn.benchmark = b
+            with torch.backends.cudnn.flags(benchmark=benchmark):
+                conv = torch.nn.Conv2d(256, 256, kernel_size=3, padding=1).type(dtype)
+                for size in sizes:
+                    x = torch.randn(size).type(dtype)
+                    out = conv(Variable(x, requires_grad=True))
+                    out.backward(torch.ones(out.size()).type(dtype))
+
+        run_test(benchmark=False)
+        run_test(benchmark=True)
 
     def test_conv_modules_raise_error_on_incorrect_input_size(self):
         modules = [nn.Conv1d(3, 8, 3), nn.ConvTranspose1d(3, 8, 3),
@@ -2631,9 +2624,7 @@ def check_rnn_grads(rnn1, rnn2):
         grad_output = torch.randn(seq_length, batch, hidden_size)
         hx_val = torch.randn(num_layers, batch, hidden_size)
         grad_hy = torch.randn(num_layers, batch, hidden_size)
-        prev = torch.backends.cudnn.enabled
-        try:
-            torch.backends.cudnn.enabled = False
+        with torch.backends.cudnn.flags(enabled=False):
             for module in (nn.GRU, nn.LSTM):
                 for bias in (True, False):
                     rnn = module(input_size, hidden_size, num_layers, bias=bias)
@@ -2676,8 +2667,6 @@ def check_rnn_grads(rnn1, rnn2):
                         self.assertEqual(hx[1].grad.data, hx_cuda[1].grad.data)
                     else:
                         self.assertEqual(hx.grad.data, hx_cuda.grad.data)
-        finally:
-            torch.backends.cudnn.enabled = prev
 
     def test_rnn_args_check(self):
         input_size = 3
@@ -2777,11 +2766,8 @@ def test_rnn_retain_variables(self):
 
     @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
     def test_rnn_retain_variables_cuda(self):
-        try:
-            torch.backends.cudnn.enabled = False
+        with torch.backends.cudnn.flags(enabled=False):
             self._test_rnn_retain_variables(torch.cuda.FloatTensor)
-        finally:
-            torch.backends.cudnn.enabled = True
         self._test_rnn_retain_variables(torch.cuda.FloatTensor)
 
     def _test_RNN_cpu_vs_cudnn(self, dropout):