Repository: pytorch/pytorch

Issue:
flatten_parameters() in RNNs for contiguous memory
Since the new pytorch version, I have a warning when I [encode sentence with a BiLSTM-max](https://github.com/facebookresearch/InferSent/blob/master/encoder/models.py#L49) that says:


```
models.py:53: UserWarning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greately increasing memory usage. To compact weights again call flatten_parameters().
sent_output = self.enc_lstm(sent_packed)[0]  # seqlen x batch x 2*nhid
```


When I do `model.enc_lstm.flatten_parameters()` as suggested, and then try to encode again, I get the following error though:


```
/home/aconneau/python/InferSent/encoder/models.py in forward(self, sent_tuple)
     51         # Handling padding in Recurrent Networks
     52         sent_packed = nn.utils.rnn.pack_padded_sequence(sent, sent_len)
---> 53         sent_output = self.enc_lstm(sent_packed)[0]  # seqlen x batch x 2*nhid
     54         sent_output = nn.utils.rnn.pad_packed_sequence(sent_output)[0]
     55 

/home/aconneau/anaconda/lib/python2.7/site-packages/torch/nn/modules/module.pyc in __call__(self, *input, **kwargs)
    222         for hook in self._forward_pre_hooks.values():
    223             hook(self, input)
--> 224         result = self.forward(*input, **kwargs)
    225         for hook in self._forward_hooks.values():
    226             hook_result = hook(self, input, result)

/home/aconneau/anaconda/lib/python2.7/site-packages/torch/nn/modules/rnn.pyc in forward(self, input, hx)
    143         if has_flat_weights:
    144             first_data = next(self.parameters()).data
--> 145             assert first_data.storage().size() == self._param_buf_size
    146             flat_weight = first_data.new().set_(first_data.storage(), 0, torch.Size([self._param_buf_size]))
    147         else:

/home/aconneau/anaconda/lib/python2.7/site-packages/torch/nn/modules/module.pyc in __getattr__(self, name)
    260                 return modules[name]
    261         raise AttributeError("'{}' object has no attribute '{}'".format(
--> 262             type(self).__name__, name))
    263 
    264     def __setattr__(self, name, value):

AttributeError: 'LSTM' object has no attribute '_param_buf_size'
```

Any idea?


Discussion:
@aconneau, do you have code the reproduces the issue?

Resolving patch:
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -32,7 +32,6 @@ def __init__(self, mode, input_size, hidden_size,
             gate_size = hidden_size
 
         self._all_weights = []
-        self._param_buf_size = 0
         for layer in range(num_layers):
             for direction in range(num_directions):
                 layer_input_size = input_size if layer == 0 else hidden_size * num_directions
@@ -53,8 +52,6 @@ def __init__(self, mode, input_size, hidden_size,
                     setattr(self, name, param)
                 self._all_weights.append(param_names)
 
-                self._param_buf_size += sum(p.numel() for p in layer_params)
-
         self.flatten_parameters()
         self.reset_parameters()
 
@@ -69,47 +66,48 @@ def flatten_parameters(self):
             self._data_ptrs = []
             return
 
-        # This is quite ugly, but it allows us to reuse the cuDNN code without larger
-        # modifications. It's really a low-level API that doesn't belong in here, but
-        # let's make this exception.
-        from torch.backends.cudnn import rnn
-        from torch.backends import cudnn
-        from torch.nn._functions.rnn import CudnnRNN
-        handle = cudnn.get_handle()
-        with warnings.catch_warnings(record=True):
-            fn = CudnnRNN(
-                self.mode,
-                self.input_size,
-                self.hidden_size,
-                num_layers=self.num_layers,
-                batch_first=self.batch_first,
-                dropout=self.dropout,
-                train=self.training,
-                bidirectional=self.bidirectional,
-                dropout_state=self.dropout_state,
-            )
-
-        # Initialize descriptors
-        fn.datatype = cudnn._typemap[any_param.type()]
-        fn.x_descs = cudnn.descriptor(any_param.new(1, self.input_size), 1)
-        fn.rnn_desc = rnn.init_rnn_descriptor(fn, handle)
-
-        # Allocate buffer to hold the weights
-        num_weights = rnn.get_num_weights(handle, fn.rnn_desc, fn.x_descs[0], fn.datatype)
-        fn.weight_buf = any_param.new(num_weights).zero_()
-        fn.w_desc = rnn.init_weight_descriptor(fn, fn.weight_buf)
-
-        # Slice off views into weight_buf
-        params = rnn.get_parameters(fn, handle, fn.weight_buf)
-        all_weights = [[p.data for p in l] for l in self.all_weights]
-
-        # Copy weights and update their storage
-        rnn._copyParams(all_weights, params)
-        for orig_layer_param, new_layer_param in zip(all_weights, params):
-            for orig_param, new_param in zip(orig_layer_param, new_layer_param):
-                orig_param.set_(new_param.view_as(orig_param))
-
-        self._data_ptrs = list(p.data.data_ptr() for p in self.parameters())
+        with torch.cuda.device_of(any_param):
+            # This is quite ugly, but it allows us to reuse the cuDNN code without larger
+            # modifications. It's really a low-level API that doesn't belong in here, but
+            # let's make this exception.
+            from torch.backends.cudnn import rnn
+            from torch.backends import cudnn
+            from torch.nn._functions.rnn import CudnnRNN
+            handle = cudnn.get_handle()
+            with warnings.catch_warnings(record=True):
+                fn = CudnnRNN(
+                    self.mode,
+                    self.input_size,
+                    self.hidden_size,
+                    num_layers=self.num_layers,
+                    batch_first=self.batch_first,
+                    dropout=self.dropout,
+                    train=self.training,
+                    bidirectional=self.bidirectional,
+                    dropout_state=self.dropout_state,
+                )
+
+            # Initialize descriptors
+            fn.datatype = cudnn._typemap[any_param.type()]
+            fn.x_descs = cudnn.descriptor(any_param.new(1, self.input_size), 1)
+            fn.rnn_desc = rnn.init_rnn_descriptor(fn, handle)
+
+            # Allocate buffer to hold the weights
+            self._param_buf_size = rnn.get_num_weights(handle, fn.rnn_desc, fn.x_descs[0], fn.datatype)
+            fn.weight_buf = any_param.new(self._param_buf_size).zero_()
+            fn.w_desc = rnn.init_weight_descriptor(fn, fn.weight_buf)
+
+            # Slice off views into weight_buf
+            params = rnn.get_parameters(fn, handle, fn.weight_buf)
+            all_weights = [[p.data for p in l] for l in self.all_weights]
+
+            # Copy weights and update their storage
+            rnn._copyParams(all_weights, params)
+            for orig_layer_param, new_layer_param in zip(all_weights, params):
+                for orig_param, new_param in zip(orig_layer_param, new_layer_param):
+                    orig_param.set_(new_param.view_as(orig_param))
+
+            self._data_ptrs = list(p.data.data_ptr() for p in self.parameters())
 
     def _apply(self, fn):
         ret = super(RNNBase, self)._apply(fn)


Test patch: