Handling Sequence Padding and Packing in PyTorch for RNNs

Implementation of Sequence Padding and Sequence Packing

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.

Python3

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.

Python3

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.

Python3

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().

Python3

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.

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.

Table of Content