How to make an RNN model in PyTorch that has a custom hidden layer(s) and that is compatible with PackedSequence

Question

I want to make an RNN that has for example more hidden layers or layer normalization.

I know that is it possible to make a custom RNN by subclassing nn.module, but with this approach is it not possible to do efficient batch processing with a PackedSequence object (with variable length sequences) the same way and with the same efficiency as torch.nn.RNN.

I thought maybe the solution could be to subclass nn.RNN, but I don't know how to do that.

Answer 1

Assuming you're using python it is possible to do (relatively) efficient batch processing with a PackedSequence object, here is some example code;

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class CustomRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, bidirectional=False):
        super(CustomRNN, self).__init__()
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, bidirectional=bidirectional, batch_first=True)
        self.layer_norm = nn.LayerNorm(hidden_size * 2 if bidirectional else hidden_size)

    def forward(self, x, lengths):
        packed_seq = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
        output, hidden = self.rnn(packed_seq)
        output, _ = pad_packed_sequence(output, batch_first=True)
        output = self.layer_norm(output)
        return output, hidden

Here, CustomRNN takes in the input_size, hidden_size, num_layers, and bidirectional parameters just like nn.RNN. In the forward method, the input sequence x and corresponding lengths are first packed into a PackedSequence object using pack_padded_sequence. The packed sequence is then passed through the RNN and the output is obtained. The output is then unpacked using pad_packed_sequence and layer normalization is applied to the output using nn.LayerNorm. Finally, the normalized output and hidden state are returned.

With this implementation, you can efficiently process variable-length sequences using a PackedSequence object while also incorporating layer normalization into the RNN.