Skip to content

Instantly share code, notes, and snippets.

@brunodoamaral

brunodoamaral/pytorch_lstm_101.py

Created March 1, 2018 19:44
Show Gist options
  • Save brunodoamaral/b0ca6f57a58bf4b9dd9c69b35f32b3d6 to your computer and use it in GitHub Desktop.
Save brunodoamaral/b0ca6f57a58bf4b9dd9c69b35f32b3d6 to your computer and use it in GitHub Desktop.
Download ZIP
Pytorch LSTM 101
Raw
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment