How to handle sequence padding and packing in PyTorch for RNNs?

There are many dataset that have sequences with variable lengths and recurrent neural networks (RNNs) require fixed-length inputs. To address this challenge, sequence padding and packing techniques are used, particularly in PyTorch, a popular deep learning framework. The article demonstrates how sequence padding ensures uniformity in sequence lengths by adding zeros to shorter sequences, while sequence packing compresses padded sequences for efficient processing in RNNs.

Sequence Padding and Packing for RNNs

Training Recurrent Neural Networks (RNNs) can be tricky when dealing with sequences of different lengths. Imagine we have a batch of 8 sequences where their lengths are: 6, 5, 4, 7, 2, 3, 8, and 7.

This is where padding comes and pad all sequences to the maximum length (8 in this case) with meaningless values. This creates an 8x8 matrix for computations, even though some sequences are shorter. This wastes processing power because we perform unnecessary calculations (64 computations instead of the actual 45 needed).

For this , packing plays an important role as It packs the sequences into a data structure that preserves their original lengths before padding. By doing so, the RNN model can process only the non-padded portions of each sequence, effectively reducing the computational overhead.

Implementation of Sequence Padding and Sequence Packing

• The code imports necessary modules from PyTorch: torch and torch.nn.utils.rnn defining a list of example sequences with variable lengths (sequences). Each sequence in the list is converted to a PyTorch tensor using a list comprehension (sequences_tensor).
• The pad_sequence function from torch.nn.utils.rnn is used to pad the sequences to the maximum length with zeros, ensuring that all sequences have the same length. The batch_first=True argument specifies that the batch dimension should be the first dimension in the resulting tensor.
• Then, the code calculates the actual lengths of sequences, and finally demonstrates how to pack sequences using pack_padded_sequence().

import torch
import torch.nn.utils.rnn as rnn_utils
# Define sequences
sequences = [  
    [1, 2, 3],
    [4, 5],
    [6, 7, 8, 9],
    [10]
]
sequences_tensor = [torch.tensor(seq) for seq in sequences] # Convert sequences to PyTorch tensors

# Padding
padded_sequences = rnn_utils.pad_sequence(sequences_tensor, batch_first=True)
print("Padded sequences:","\n",padded_sequences)

# Packing 
sequence_lengths = torch.tensor([len(seq) for seq in sequences]) # Calculating actual lengths of sequences
# Pack padded sequences
packed_sequences = rnn_utils.pack_padded_sequence(padded_sequences, sequence_lengths, batch_first=True, enforce_sorted=False)
print("\nPacked sequences:",packed_sequences)

Output:

Padded sequences: 
 tensor([[ 1,  2,  3,  0],
        [ 4,  5,  0,  0],
        [ 6,  7,  8,  9],
        [10,  0,  0,  0]])

Packed sequences: PackedSequence(data=tensor([ 6,  1,  4, 10,  7,  2,  5,  8,  3,  9]), batch_sizes=tensor([4, 3, 2, 1]), sorted_indices=tensor([2, 0, 1, 3]), unsorted_indices=tensor([1, 2, 0, 3]))

The output consists of a 2-dimensional PyTorch tensor, representing the padded sequences. Each row in the tensor corresponds to a sequence, and columns represent elements within each sequence. For example,

• First Sequence (Row 1):
  • Original sequence: [1, 2, 3]
  • Padded sequence: [1, 2, 3, 0]
  • The original sequence had three elements, so it was padded with a zero to match the length of the longest sequence in the batch (which is four).

• Second Sequence (Row 2):
  • Original sequence: [4, 5]
  • Padded sequence: [4, 5, 0, 0]
  • The original sequence had two elements, so it was padded with two zeros to match the length of the longest sequence in the batch.

Same is done for all the rows.

In the packed sequence:

• data: contains the flattened non-padded elements from the padded sequences.
• batch_sizes: indicates how many elements are present at each time step, reflecting the varying sequence lengths within the batch.

This packed sequence is feed into your recurrent neural network (RNN) model during training, allowing it to efficiently process variable-length sequences.

Handling Sequence Padding and Packing in PyTorch for RNNs

This code implements a basic RNN model using PyTorch's nn.Module class. for sequence processing tasks, while handling variable-length input sequences using sequence packing and unpacking techniques.

The forward method takes input sequences (text) and their lengths (text_lengths). Inside the forward method, sequence packing is performed.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Define RNN Model
class RNN(nn.Module):
    def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, embedding_dim)
        self.rnn = nn.RNN(embedding_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, text, text_lengths):
        embedded = self.embedding(text)
        packed_embedded = pack_padded_sequence(embedded, text_lengths.cpu(), batch_first=True, enforce_sorted=False)
        packed_output, _ = self.rnn(packed_embedded)
        output, output_lengths = pad_packed_sequence(packed_output, batch_first=True)
        return self.fc(output[:, -1, :])

Pad Sequences

• The example sequences are padded to the maximum length using torch.nn.utils.rnn.pad_sequence(). This is necessary for processing sequences in batches.
• The lengths of the original sequences are computed and stored in a tensor (sequence_lengths). This is necessary for sequence packing.

sequences = [
    torch.tensor([1, 2, 3]),
    torch.tensor([4, 5]),
    torch.tensor([6, 7, 8, 9]),
    torch.tensor([10])
]
padded_sequences = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)
sequence_lengths = torch.tensor([len(seq) for seq in sequences])

Instantiate the Model

An instance of the RNN model is created with example dimensions for input, embedding, hidden, and output layers to determine the architecture and behavior of the model.

INPUT_DIM = 11
EMBEDDING_DIM = 100
HIDDEN_DIM = 256
OUTPUT_DIM = 1
model = RNN(INPUT_DIM, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM)

Forward Pass with Packed Sequences

The model is called with the padded sequences (padded_sequences) and their lengths (sequence_lengths). Inside the model's forward method, sequence packing is performed using pack_padded_sequence().

outputs = model(padded_sequences, sequence_lengths)
print("Output shape:", outputs.shape)

Output:

Output shape: torch.Size([4, 1])

The output shape torch.Size([4, 1]) indicates that the model produced a tensor with 4 rows and 1 column.

• Rows: The first dimension of the tensor corresponds to the batch size. In this case, there are 4 sequences in the batch.
• Columns: The second dimension corresponds to the output dimension of the model. Since OUTPUT_DIM = 1, each sequence produces a single output value.

Conclusion

In conclusion, sequence padding ensures uniformity in sequence lengths by adding zeros to shorter sequences, while sequence packing compresses padded sequences for efficient processing in recurrent neural networks (RNNs) using PyTorch.