Pytorch LSTM 101

pytorch_lstm_101.py

from torch.autograd import Variable
import torch.nn as nn
import torch

# LSTM
hidden_size = 20
num_layers = 2
num_directions = 1

# Input
input_size = 10
seq_len = 5
batch = 3

# Instantiate LSTM
rnn = nn.LSTM(input_size, hidden_size, num_layers, bidirectional=(num_directions == 2))

# Input simulation
input = Variable(torch.randn(seq_len, batch, input_size))

# Hidden states h and c
h_0 = Variable(torch.randn(num_layers * num_directions, batch, hidden_size))
c_0 = Variable(torch.randn(num_layers * num_directions, batch, hidden_size))

# Capture output
output, (h_n, c_n) = rnn(input, (h_0, c_0))

# Shape assertions
assert output.shape == (seq_len, batch, hidden_size * num_directions)
assert h_n.shape == (num_layers * num_directions, batch, hidden_size)
assert c_n.shape == (num_layers * num_directions, batch, hidden_size)

# For debug
output.shape, h_n.shape, c_n.shape