Repository: pytorch/pytorch

Issue:
Bind cuDNN dropout initialization from ATen
In `torch/backends/cudnn/rnn.py` we have the last vestiges of the Python-side cuDNN bindings:

```
class Unserializable(object):

    def __init__(self, inner):
        self.inner = inner

    def get(self):
        return self.inner

    def __getstate__(self):
        # Note: can't return {}, because python2 won't call __setstate__
        # if the value evaluates to False
        return "<unserializable>"

    def __setstate__(self, state):
        self.inner = None


def init_dropout_descriptor(handle, dropout, train, dropout_seed, dropout_state):
    dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
    dropout_p = dropout if train else 0
    if (dropout_desc_name not in dropout_state) or (dropout_state[dropout_desc_name].get() is None):
        dropout_state[dropout_desc_name] = Unserializable(
            cudnn.DropoutDescriptor(handle, dropout_p, dropout_seed)
        )
    dropout_desc = dropout_state[dropout_desc_name].get()
    dropout_desc.set_dropout(dropout_p, dropout_seed)
    return dropout_desc


def get_dropout_state(fn, handle):
    dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
    dropout_p = fn.dropout if fn.train else 0
    dropout_desc = fn.dropout_state[dropout_desc_name].get()
    return dropout_desc.state
```

We'd like to get rid of this code entirely, and that means we need to allocate the tensor state in `dropout_desc` from ATen.

The code to do this already exists in `aten/src/ATen/cudnn/Descriptors.h`; it's just a matter of taking dropout parameters and returning an initialized dropout state buffer, which can be passed to subsequent ATen calls.


Discussion:
We can't remove the `Unserializable` class or people won't be able to load older models.
This is a little annoying to do as a native function as we have never done the zero Tensor arguments returns a Tensor pattern before.

Resolving patch:
diff --git a/torch/backends/cudnn/rnn.py b/torch/backends/cudnn/rnn.py
--- a/torch/backends/cudnn/rnn.py
+++ b/torch/backends/cudnn/rnn.py
@@ -1,8 +1,5 @@
 import torch.cuda
 import torch.backends.cudnn as cudnn
-from torch.backends.cudnn import check_error
-import ctypes
-from torch.autograd import Variable
 
 
 def get_cudnn_mode(mode):
@@ -18,6 +15,9 @@ def get_cudnn_mode(mode):
         raise Exception("Unknown mode: {}".format(mode))
 
 
+# NB: We don't actually need this class anymore (in fact, we could serialize the
+# dropout state for even better reproducibility), but it is kept for backwards
+# compatibility for old models.
 class Unserializable(object):
 
     def __init__(self, inner):
@@ -35,20 +35,13 @@ def __setstate__(self, state):
         self.inner = None
 
 
-def init_dropout_descriptor(handle, dropout, train, dropout_seed, dropout_state):
+def init_dropout_state(ty, dropout, train, dropout_seed, dropout_state):
     dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
     dropout_p = dropout if train else 0
     if (dropout_desc_name not in dropout_state) or (dropout_state[dropout_desc_name].get() is None):
         dropout_state[dropout_desc_name] = Unserializable(
-            cudnn.DropoutDescriptor(handle, dropout_p, dropout_seed)
+            torch._C._VariableFunctions._cudnn_init_dropout_state(ty, dropout_p, train, dropout_seed)
+            if dropout_p != 0 else None
         )
-    dropout_desc = dropout_state[dropout_desc_name].get()
-    dropout_desc.set_dropout(dropout_p, dropout_seed)
-    return dropout_desc
-
-
-def get_dropout_state(fn, handle):
-    dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
-    dropout_p = fn.dropout if fn.train else 0
-    dropout_desc = fn.dropout_state[dropout_desc_name].get()
-    return dropout_desc.state
+    dropout_ts = dropout_state[dropout_desc_name].get()
+    return dropout_ts
diff --git a/torch/nn/_functions/rnn.py b/torch/nn/_functions/rnn.py
--- a/torch/nn/_functions/rnn.py
+++ b/torch/nn/_functions/rnn.py
@@ -255,7 +255,8 @@ def CudnnRNN(mode, input_size, hidden_size, num_layers=1,
     if dropout_state is None:
         dropout_state = {}
     mode = cudnn.rnn.get_cudnn_mode(mode)
-    dropout_seed = torch.IntTensor(1).random_()[0]
+    # TODO: This is really goofy way of using the Torch RNG to get a random number
+    dropout_seed = int(torch.IntTensor(1).random_())
     if flat_weight is None:
         warnings.warn("RNN module weights are not part of single contiguous "
                       "chunk of memory. This means they need to be compacted "
@@ -269,7 +270,7 @@ def forward(input, weight, hx, batch_sizes):
             cx = None
 
         handle = cudnn.get_handle()
-        dropout_desc = cudnn.rnn.init_dropout_descriptor(handle, dropout, train, dropout_seed, dropout_state)
+        dropout_ts = cudnn.rnn.init_dropout_state(torch.cuda.uint8, dropout, train, dropout_seed, dropout_state)
 
         weight_arr = list(itertools.chain.from_iterable(weight))
         weight_stride0 = len(weight[0])
@@ -281,7 +282,7 @@ def forward(input, weight, hx, batch_sizes):
             mode, hidden_size, num_layers,
             batch_first, dropout, train, bool(bidirectional),
             list(batch_sizes.data) if variable_length else (),
-            Variable(dropout_desc.state) if dropout_desc.state is not None else None)
+            dropout_ts)
 
         if cx is not None:
             return (output, (hy, cy))


Test patch:
diff --git a/aten/src/ATen/test/cudnn_test.cpp b/aten/src/ATen/test/cudnn_test.cpp
--- a/aten/src/ATen/test/cudnn_test.cpp
+++ b/aten/src/ATen/test/cudnn_test.cpp
@@ -10,7 +10,7 @@ int main() {
 #if CUDNN_VERSION < 7000
   auto handle = getCudnnHandle();
   DropoutDescriptor desc1, desc2;
-  desc1.initialize_rng(handle, 0.5, 42);
+  desc1.initialize_rng(at::CUDA(kByte), handle, 0.5, 42);
   desc2.set(handle, 0.5, desc1.state);
 
   ASSERT(desc1.desc()->dropout == desc2.desc()->dropout);