Repository: speechbrain/speechbrain

Issue:
RNNs do not handle padding correctly
This is an issue discovered with the help of @JasonSWFu , the behavior of RNNs is different based on the batch size, presumably due to padding differences. I'm pretty sure the intended way to handle variable-length batches in PyTorch is to use `torch.nn.utils.rnn.pack_padded_sequence` on tensors before passing them to RNNs, so that padding is ignored.

The two approaches to solving it that came to mind:

1. Add `pack_padded_sequence` to each RNN's `forward()` method. Unfortunately this function requires the length, so this approach also requires updating every call to an RNN module, including those inside of `Sequential` containers.
2. Warn users when an RNN gets a tensor rather than a `PackedSequence` object, so they can manually add the `pad_packed_sequence` call. Unfortunately, this approach also doesn't handle `Sequential` containers.

This seems common enough that perhaps some sort of intervention in `Sequential` is appropriate, like always passing the length if there are two parameters to the forward method.

I think this issue is pretty important, how should we proceed with fixing it?


Discussion:

Resolving patch:
diff --git a/speechbrain/nnet/RNN.py b/speechbrain/nnet/RNN.py
--- a/speechbrain/nnet/RNN.py
+++ b/speechbrain/nnet/RNN.py
@@ -21,6 +21,36 @@
 logger = logging.getLogger(__name__)
 
 
+def pack_padded_sequence(inputs, lengths):
+    """Returns packed speechbrain-formatted tensors.
+
+    Arguments
+    ---------
+    inputs : torch.Tensor
+        The sequences to pack.
+    lengths : torch.Tensor
+        The length of each sequence.
+    """
+    lengths = (lengths * inputs.size(1)).cpu()
+    return torch.nn.utils.rnn.pack_padded_sequence(
+        inputs, lengths, batch_first=True, enforce_sorted=False
+    )
+
+
+def pad_packed_sequence(inputs):
+    """Returns speechbrain-formatted tensor from packed sequences.
+
+    Arguments
+    ---------
+    inputs : torch.nn.utils.rnn.PackedSequence
+        An input set of sequences to convert to a tensor.
+    """
+    outputs, lengths = torch.nn.utils.rnn.pad_packed_sequence(
+        inputs, batch_first=True
+    )
+    return outputs
+
+
 class RNN(torch.nn.Module):
     """This function implements a vanilla RNN.
 
@@ -99,12 +129,17 @@ def __init__(
         if re_init:
             rnn_init(self.rnn)
 
-    def forward(self, x, hx=None):
+    def forward(self, x, hx=None, lengths=None):
         """Returns the output of the vanilla RNN.
 
         Arguments
         ---------
         x : torch.Tensor
+            Input tensor.
+        hx : torch.Tensor
+            Starting hidden state.
+        lengths : torch.Tensor
+            Relative lengths of the input signals.
         """
         # Reshaping input tensors for 4d inputs
         if self.reshape:
@@ -114,12 +149,20 @@ def forward(self, x, hx=None):
         # Flatten params for data parallel
         self.rnn.flatten_parameters()
 
+        # Pack sequence for proper RNN handling of padding
+        if lengths is not None:
+            x = pack_padded_sequence(x, lengths)
+
         # Support custom inital state
         if hx is not None:
             output, hn = self.rnn(x, hx=hx)
         else:
             output, hn = self.rnn(x)
 
+        # Unpack the packed sequence
+        if lengths is not None:
+            output = pad_packed_sequence(output)
+
         return output, hn
 
 
@@ -197,13 +240,17 @@ def __init__(
         if re_init:
             rnn_init(self.rnn)
 
-    def forward(self, x, hx=None):
+    def forward(self, x, hx=None, lengths=None):
         """Returns the output of the LSTM.
 
         Arguments
         ---------
         x : torch.Tensor
             Input tensor.
+        hx : torch.Tensor
+            Starting hidden state.
+        lengths : torch.Tensor
+            Relative length of the input signals.
         """
         # Reshaping input tensors for 4d inputs
         if self.reshape:
@@ -213,12 +260,20 @@ def forward(self, x, hx=None):
         # Flatten params for data parallel
         self.rnn.flatten_parameters()
 
+        # Pack sequence for proper RNN handling of padding
+        if lengths is not None:
+            x = pack_padded_sequence(x, lengths)
+
         # Support custom inital state
         if hx is not None:
             output, hn = self.rnn(x, hx=hx)
         else:
             output, hn = self.rnn(x)
 
+        # Unpack the packed sequence
+        if lengths is not None:
+            output = pad_packed_sequence(output)
+
         return output, hn
 
 
@@ -296,12 +351,17 @@ def __init__(
         if re_init:
             rnn_init(self.rnn)
 
-    def forward(self, x, hx=None):
+    def forward(self, x, hx=None, lengths=None):
         """Returns the output of the GRU.
 
         Arguments
         ---------
         x : torch.Tensor
+            Input tensor.
+        hx : torch.Tensor
+            Starting hidden state.
+        lengths : torch.Tensor
+            Relative length of the input signals.
         """
         # Reshaping input tensors for 4d inputs
         if self.reshape:
@@ -311,12 +371,20 @@ def forward(self, x, hx=None):
         # Flatten params for data parallel
         self.rnn.flatten_parameters()
 
+        # Pack sequence for proper RNN handling of padding
+        if lengths is not None:
+            x = pack_padded_sequence(x, lengths)
+
         # Support custom inital state
         if hx is not None:
             output, hn = self.rnn(x, hx=hx)
         else:
             output, hn = self.rnn(x)
 
+        # Unpack the packed sequence
+        if lengths is not None:
+            output = pad_packed_sequence(output)
+
         return output, hn
 
 
diff --git a/speechbrain/nnet/containers.py b/speechbrain/nnet/containers.py
--- a/speechbrain/nnet/containers.py
+++ b/speechbrain/nnet/containers.py
@@ -49,6 +49,9 @@ class Sequential(torch.nn.ModuleDict):
     def __init__(self, *layers, input_shape=None, **named_layers):
         super().__init__()
 
+        # Keep track of what layers need "lengths" passed
+        self.length_layers = []
+
         # Replace None dimensions with arbitrary value
         self.input_shape = input_shape
         if input_shape and None in input_shape:
@@ -121,8 +124,8 @@ def forward(self, x):
 
         Arguments
         ---------
-        x : tensor
-            the input tensor to run through the network.
+        x : torch.Tensor
+            The input tensor to run through the network.
         """
         for layer in self.values():
             x = layer(x)
@@ -132,6 +135,53 @@ def forward(self, x):
         return x
 
 
+class LengthCapableSequential(Sequential):
+    """Sequential model that can take ``lengths`` in the forward method.
+
+    This is useful for Sequential models that include RNNs where it is
+    important to avoid padding, or for some feature normalization layers.
+
+    Unfortunately, this module is not jit-able because the compiler doesn't
+    know ahead of time if the length will be passed, and some layers don't
+    accept the length parameter.
+    """
+
+    def forward(self, x, lengths=None):
+        """Applies layers in sequence, passing only the first element of tuples.
+
+        In addition, forward the ``lengths`` argument to all layers that accept
+        a ``lengths`` argument in their ``forward()`` method (e.g. RNNs).
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            The input tensor to run through the network.
+        lengths : torch.Tensor
+            The relative lengths of each signal in the tensor.
+        """
+        for layer in self.values():
+            if lengths_arg_exists(layer):
+                x = layer(x, lengths=lengths)
+            else:
+                x = layer(x)
+            if isinstance(x, tuple):
+                x = x[0]
+
+        return x
+
+
+def lengths_arg_exists(module):
+    """Returns True if module takes ``lengths`` keyword argument.
+
+    Arguments
+    ---------
+    module : torch.nn.Module
+        Torch module to use for determination.
+    """
+    spec = inspect.getfullargspec(module.forward)
+    return "lengths" in spec.args + spec.kwonlyargs
+
+
 class ModuleList(torch.nn.Module):
     """This class implements a wrapper to torch.nn.ModuleList with a forward()
     method to forward all the layers sequentially.


Test patch: