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.


1 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.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.