Repository: pytorch/pytorch

Issue:
[feature request] make pad_packed_sequence work in DataParallel
Currently users can not do `pack -> RNN -> unpack` in a module wrapped in `DataParallel` because the unpack operation (`pad_packed_sequence`) will only pad up to the longest input it sees, i.e., the longest on that particular device. Then the code breaks when it tries to gather the results into a single tensor afterwards. I propose to add a `total_length` option to `pad_packed_sequence` so that it can beyond the longest input. Then the following pattern should work with DataParallel:
```python
def forward(self, input, input_lengths):
    total_length = input.size(-1)
    packed_input = nn.utils.rnn.pack_padded_sequence(input, input_lengths)
    packed_output = self.lstm(packed_input)
    output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, total_length=total_length)
    return output
```

The current suggested workaround is to manually add the padding. However, it becomes tricky when one wants to do ONNX export.

Relevant post: https://discuss.pytorch.org/t/question-about-packed-rnn-with-dataparallel/2738

cc @goldsborough 


Discussion:
@apaszke do you think that this is reasonable? :)

Resolving patch:
diff --git a/torch/nn/parallel/data_parallel.py b/torch/nn/parallel/data_parallel.py
--- a/torch/nn/parallel/data_parallel.py
+++ b/torch/nn/parallel/data_parallel.py
@@ -52,8 +52,22 @@ class DataParallel(Module):
 
     .. warning::
         Forward and backward hooks defined on :attr:`module` and its submodules
-        won't be invoked anymore, unless the hooks are initialized in the
-        :meth:`forward` method.
+        will be invoked ``len(device_ids)`` times, each with inputs located on
+        a particular device. Particularly, the hooks are only guaranteed to be
+        executed in correct order with respect to operations on corresponding
+        devices. For example, it is not guaranteed that hooks set via
+        :meth:`~torch.nn.Module.register_forward_pre_hook` be executed before
+        `all` ``len(device_ids)`` :meth:`~torch.nn.Module.forward` calls, but
+        that each such hook be executed before the corresponding
+        :meth:`~torch.nn.Module.forward` call of that device.
+
+    .. note::
+        There is a subtlety in using the
+        ``pack sequence -> recurrent network -> unpack sequence`` pattern in a
+        :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.
+        See :ref:`this FAQ section <pack-rnn-unpack-with-data-parallelism>` for
+        details.
+
 
     Args:
         module: module to be parallelized
diff --git a/torch/nn/utils/rnn.py b/torch/nn/utils/rnn.py
--- a/torch/nn/utils/rnn.py
+++ b/torch/nn/utils/rnn.py
@@ -11,7 +11,7 @@
 
 
 class PackedSequence(PackedSequence_):
-    r"""Holds the data and list of batch_sizes of a packed sequence.
+    r"""Holds the data and list of :attr:`batch_sizes` of a packed sequence.
 
     All RNN modules accept packed sequences as inputs.
 
@@ -21,8 +21,9 @@ class PackedSequence(PackedSequence_):
 
         Batch sizes represent the number elements at each sequence step in
         the batch, not the varying sequence lengths passed to
-        :func:`pack_padded_sequence`.  For instance, given data  ``abc`` and `d`
-        the ``PackedSequence`` would be ``adbc`` with ``batch_sizes=[2,1,1]``.
+        :func:`pack_padded_sequence`.  For instance, given data  ``abc`` and `x`
+        the :class:`PackedSequence` would contain data ``axbc`` with
+        ``batch_sizes=[2,1,1]``.
 
     Attributes:
         data (Variable): Variable containing packed sequence
@@ -136,7 +137,9 @@ def pack_padded_sequence(input, lengths, batch_first=False):
     return PackedSequence(data, batch_sizes)
 
 
-def _symbolic_pad_packed_sequence(g, input, batch_first=False, padding_value=0.0):
+def _symbolic_pad_packed_sequence(g, input, batch_first=False, padding_value=0.0, total_length=None):
+    if total_length is not None:
+        raise ValueError("_symbolic_pad_packed_sequence only supports total_length=None")
     # See comment on _symbolic_pack_padded_sequence
     data, lengths = g.op("prim::PadPacked", input.data, input.batch_sizes, outputs=2)
     if batch_first:
@@ -145,7 +148,7 @@ def _symbolic_pad_packed_sequence(g, input, batch_first=False, padding_value=0.0
 
 
 @torch.onnx.symbolic_override_packed_sequence_based(_symbolic_pad_packed_sequence)
-def pad_packed_sequence(sequence, batch_first=False, padding_value=0):
+def pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None):
     r"""Pads a packed batch of variable length sequences.
 
     It is an inverse operation to :func:`pack_padded_sequence`.
@@ -156,11 +159,22 @@ def pad_packed_sequence(sequence, batch_first=False, padding_value=0):
 
     Batch elements will be ordered decreasingly by their length.
 
+    .. note::
+        :attr:`total_length` is useful to implement the
+        ``pack sequence -> recurrent network -> unpack sequence`` pattern in a
+        :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.
+        See :ref:`this FAQ section <pack-rnn-unpack-with-data-parallelism>` for
+        details.
+
     Arguments:
         sequence (PackedSequence): batch to pad
         batch_first (bool, optional): if ``True``, the output will be in ``B x T x *``
             format.
         padding_value (float, optional): values for padded elements.
+        total_length (int, optional): if not ``None``, the output will be padded to
+            have length :attr:`total_length`. This method will throw :class:`ValueError`
+            if :attr:`total_length` is less than the max sequence length in
+            :attr:`sequence`.
 
     Returns:
         Tuple of Variable containing the padded sequence, and Variable
@@ -169,7 +183,15 @@ def pad_packed_sequence(sequence, batch_first=False, padding_value=0):
     """
     var_data, batch_sizes = sequence
     max_batch_size = int(batch_sizes[0])
-    output = var_data.data.new(len(batch_sizes), max_batch_size, *var_data.size()[1:]).fill_(padding_value)
+    max_seq_length = batch_sizes.size(0)
+    if total_length is not None:
+        if total_length < max_seq_length:
+            raise ValueError("Expected total_length to be at least the length "
+                             "of the longest sequence in input, but got "
+                             "total_length={} and max sequence length being {}"
+                             .format(total_length, max_seq_length))
+        max_seq_length = total_length
+    output = var_data.data.new(max_seq_length, max_batch_size, *var_data.size()[1:]).fill_(padding_value)
     output = Variable(output)
 
     lengths = []
diff --git a/torch/onnx/__init__.py b/torch/onnx/__init__.py
--- a/torch/onnx/__init__.py
+++ b/torch/onnx/__init__.py
@@ -124,7 +124,12 @@ def symbolic_override_first_arg_based(symbolic_fn):
 
     def might_trace(args):
         import torch
-        return torch._C._jit_is_tracing(args[0])
+        first_arg = args[0]
+        if not torch.is_tensor(first_arg):
+            raise ValueError('First argument of {} is expected to be a tensor, '
+                             'but got an object of type {}'
+                             .format(symbolic_fn.__name__, type(first_arg)))
+        return torch._C._jit_is_tracing(first_arg)
 
     return functools.partial(_symbolic_override_wrapper_maker, symbolic_fn, might_trace)
 
@@ -140,6 +145,11 @@ def symbolic_override_packed_sequence_based(symbolic_fn):
 
     def might_trace(args):
         import torch
-        return torch._C._jit_is_tracing(args[0][0])
+        first_arg = args[0]
+        if not isinstance(first_arg, torch.nn.utils.rnn.PackedSequence):
+            raise ValueError('pad_packed_sequence expects sequence to be a '
+                             'PackedSequence, but got an object of type {}'
+                             .format(type(first_arg)))
+        return torch._C._jit_is_tracing(first_arg[0])
 
     return functools.partial(_symbolic_override_wrapper_maker, symbolic_fn, might_trace)


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -128,6 +128,39 @@ def test_cuda_mask(self):
         unpacked, _ = rnn_utils.pad_packed_sequence(packed)
         self.assertEqual(unpacked.type(), cuda_type_str)
 
+    def test_total_length(self):
+        padded, lengths = self._padded_sequence(torch.FloatTensor)
+        max_length = max(lengths)
+        packed = rnn_utils.pack_padded_sequence(padded, lengths)
+        # test ValueError if total_length < max_length
+        for total_length in (-1, 0, max_length - 1):
+            for batch_first in (True, False):
+                def err_fn():
+                    rnn_utils.pad_packed_sequence(packed, batch_first=batch_first,
+                                                  total_length=total_length)
+            self.assertRaisesRegex(ValueError,
+                                   r'Expected total_length to be at least the '
+                                   r'length of the longest sequence in input',
+                                   err_fn)
+        # test that pad_packed_sequence returns results of correct length
+        for batch_first in (True, False):
+            no_extra_pad, _ = rnn_utils.pad_packed_sequence(packed, batch_first=batch_first)
+            for total_length_delta in (0, 1, 8):
+                total_length = max_length + total_length_delta
+                unpacked, lengths_out = rnn_utils.pad_packed_sequence(packed, batch_first=batch_first,
+                                                                      total_length=total_length)
+                self.assertEqual(lengths, lengths_out)
+                self.assertEqual(unpacked.size(1 if batch_first else 0), total_length)
+                if total_length_delta == 0:
+                    ref_output = no_extra_pad
+                elif batch_first:
+                    extra_pad = no_extra_pad.new_zeros(self.batch_size, total_length_delta)
+                    ref_output = torch.cat([no_extra_pad, extra_pad], 1)
+                else:
+                    extra_pad = no_extra_pad.new_zeros(total_length_delta, self.batch_size)
+                    ref_output = torch.cat([no_extra_pad, extra_pad], 0)
+                self.assertEqual(unpacked, ref_output)
+
 
 def default_tensor_type(type):
     type_str = torch.typename(type)