In [1]:
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data_utils
import operator
import numpy as np
Let us assume that we have an input $x = [x_1, x_2, ..., x_N]$ and we need to learn the mapping for some output $y = [y_1, y_2, ..., y_N]$, where $N$ is variable for each instance. In this case we can't just use a simple feed forward neural network which maps $x \rightarrow y$, as this will not work with variable length sequences. Furthermore, the number or parameters required for training such a network would be proportional to $size(x_i)*N$. This is a major memory cost. Additionally, if the sequence has some common mapping between $x_i$ and $y_i$, then we would be learning redundant weights for each pair in the sequence. This is where an RNN network is more useful. The basic idea is that each input $x_i$ is processed in a similar fashion using the same processing module and some additional context variable (which we will henseforth refer to as the hidden state). This hidden state should capture some information about the part of the sequence which has already been processed. Now at each step of the sequence we need to do the following:
The order of the above steps is not fixed and forms the basis of many RNN spin-offs. What is important, at each step, is to have a new output and a new hidden state. Sometimes, the hidden state and the outputs are the same, to make the network smaller. But the core idea remains same. Below we would like to formalize the general intuition of an RNN module.
Initialize an initial hidden state $h_{0}$ with some initial value.
At timestep n: $$ \begin{equation} h^{'}_{i} = f(x_{i},h_{i})\\ y_{i} = g(x_{i},h^{'}_{i})\\ h_{i+1} = h^{'}_{i}\\ \end{equation} $$
Here $y_{i}$ is the output and $h^{'}_{i}$ is the intermediate hidden state.
In [2]:
class Input2Hidden(nn.Module):
def __init__(self, x_dim, concat_layers=False):
"""Input2Hidden module
Args:
x_dim: input vector dimension
concat_layers: weather to concat input and hidden layers or sum them
"""
super(Input2Hidden, self).__init__()
self.concat_layers = concat_layers
input_dim = x_dim
if self.concat_layers:
input_dim = 2*x_dim
self.linear_layer = nn.Linear(input_dim, x_dim)
def forward(self, x, h):
if self.concat_layers:
cell_input = torch.cat([x,h], dim=1)
else:
cell_input = x + h
assert isinstance(cell_input, Variable)
logit = F.tanh(self.linear_layer(cell_input))
return logit
class Hidden2Output(nn.Module):
def __init__(self, x_dim, out_dim, concat_layers=False):
"""Hidden2Output module
Args:
x_dim: input vector dimension
out_dim: output vector dimension
concat_layers: weather to concat input and hidden layers or sum them
"""
super(Hidden2Output, self).__init__()
input_dim = x_dim
self.concat_layers = concat_layers
if self.concat_layers:
input_dim = 2*x_dim
self.linear_layer = nn.Linear(input_dim, out_dim)
def forward(self, x, h):
if self.concat_layers:
cell_input = torch.cat([x,h], dim=1)
else:
cell_input = x + h
assert isinstance(cell_input, Variable)
logit = F.tanh(self.linear_layer(cell_input))
return logit
class CustomRNNCell(nn.Module):
def __init__(self, i2h, h2o):
super(CustomRNNCell, self).__init__()
self.i2h = i2h
self.h2o = h2o
def forward(self, x, h):
assert isinstance(x, Variable)
assert isinstance(h, Variable)
h_prime = self.i2h(x,h)
assert isinstance(h_prime, Variable)
output = self.h2o(x,h_prime)
return output, h_prime
class Model(nn.Module):
def __init__(self, embedding, rnn_cell):
super(Model, self).__init__()
self.embedding = embedding
self.rnn_cell = rnn_cell
self.loss_function = nn.CrossEntropyLoss()
def forward(self, word_ids, hidden=None):
if hidden is None:
hidden = Variable(torch.zeros(
word_ids.data.shape[0],self.embedding.embedding_dim))
assert isinstance(hidden, Variable)
embeddings = self.embedding(word_ids)
max_seq_length = word_ids.data.shape[-1]
outputs, hidden_states = [], []
for i in range(max_seq_length):
x = embeddings[:, i, :]
assert isinstance(x, Variable)
#print("x={}\nhidden={}".format(x,hidden))
output, hidden = self.rnn_cell(x, hidden)
assert isinstance(output, Variable)
assert isinstance(hidden, Variable)
#print("output: {}, hidden: {}".format(output.data.shape, hidden.data.shape))
outputs.append(output.unsqueeze(1))
hidden_states.append(hidden.unsqueeze(1))
outputs = torch.cat(outputs, 1)
hidden_states = torch.cat(hidden_states, 1)
assert isinstance(outputs, Variable)
assert isinstance(hidden_states, Variable)
return outputs, hidden_states
def loss(self, word_ids, target_ids, hidden=None):
outputs, hidden_states = self.forward(word_ids, hidden=hidden)
outputs = outputs.view(-1, outputs.data.shape[-1])
target_ids = target_ids.view(-1)
assert isinstance(outputs, Variable)
assert isinstance(target_ids, Variable)
#print("output={}\ttargets={}".format(outputs.data.shape,target_ids.data.shape))
loss = self.loss_function(outputs, target_ids)
return loss
def predict(self, word_ids, hidden=None):
outputs, hidden_states = self.forward(word_ids, hidden=hidden)
outputs = outputs.view(-1, outputs.data.shape[-1])
max_scores, predictions = outputs.max(1)
predictions = predictions.view(*word_ids.data.shape)
#print(word_ids.data.shape, predictions.data.shape)
assert word_ids.data.shape == predictions.data.shape, "word_ids: {}, predictions: {}".format(
word_ids.data.shape, predictions.data.shape
)
return predictions
def tensors2variables(*args, requires_grad=False):
return tuple(map(lambda x: Variable(x, requires_grad=requires_grad), args))
def get_batch(tensor_types, *args, requires_grad=False):
return tuple(map(lambda t,arg: Variable(t(arg), requires_grad=requires_grad), tensor_types, args))
Let us take a simple example of using an RNN to predict the flip in bits of an $N$ bit unsigned integer. In python for an integer n represented using $N$ bits, the unsigned bitflip can be written as (~n) & ((1 << N)-1). Four our RNN each bit from the left will be $x_i$ and each flipped bit will be $y_i$. This task doesn't require any temporal dependencies but will be a good exercise to test the accuracy of RNN implementation. Theoretically, the network should learn to do this job perfectly in a few iterations. Later we will move to an example which does require the network to learn some temporal dependencies between inputs.
For our network we define $f(x,h)$ as a simple affine layer with $tanh$ activation, which takes the concatanated input $[x_i, h_{i-1}]$ and returns a new hidden state $h^{'}_{i}$. Similarly, we have $g(x_i, h^{'}_{n})$ also represented as an affine layer with $tanh$ activation, of the concatanation of its inputs $[x_i, h^{'}_{n}]$, resulting in a new output $y_{i}$. More formally, we have
$$ \begin{equation} h^{'}_{i} = f(x_{i},h_{i}) = \sigma([x_i, h_{i-1}]W_{i2h})\\ y_{i} = g(x_{i},h^{'}_{i}) = \sigma([x_i, h^{'}_{n}]W_{h2o})\\ h_{i+1} = h^{'}_{i}\\ \end{equation} $$
In [3]:
input_size=2
embedding_size=3
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
f = Input2Hidden(embedding_size, concat_layers=True)
g = Hidden2Output(embedding_size, 2, concat_layers=True)
rnn_cell = CustomRNNCell(f,g)
model = Model(embedding, rnn_cell)
In [4]:
word_ids = [[0, 1, 0, 1, 0, 1]]
target_ids = [[1, 0, 1, 0, 1, 0]]
tensor_types = (torch.LongTensor, torch.LongTensor)
word_ids_tensor, target_ids_tensor = get_batch(tensor_types, word_ids, target_ids)
print(word_ids_tensor, target_ids_tensor)
In [5]:
model.forward(torch.cat([word_ids_tensor, word_ids_tensor, word_ids_tensor], 0))[0]
Out[5]:
In [6]:
model.predict(torch.cat([word_ids_tensor, word_ids_tensor, word_ids_tensor], 0))
Out[6]:
In [7]:
loss = model.loss(word_ids_tensor, target_ids_tensor)
loss
Out[7]:
In [8]:
loss.backward()
In [9]:
model.predict(word_ids_tensor)
Out[9]:
In [10]:
def create_dataset(max_len=5):
"""Create a dataset of max_len bits and their flipped values
Args:
max_len: Maximum number of bits in the number
"""
max_val = (1<<max_len)
X, Y = [], []
for i in range(max_val):
x = "{0:0{1}b}".format(i,max_len)
y = "{0:0{1}b}".format((~i) & max_val-1,max_len)
x = tuple(map(int, x))
y = tuple(map(int, y))
X.append(x)
Y.append(y)
return X, Y
In [11]:
X, Y = create_dataset(max_len=10)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
In [12]:
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=8, shuffle=True)
In [13]:
%%time
learning_rate = 1e-4
max_epochs = 20
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % 2 != 0:
continue
loss = model.loss(*tensors2variables(X_tensors, Y_tensors))
Y_predict = model.predict(Variable(X_tensors)).data
accuracy = (Y_tensors == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
In [14]:
Y_predict = model.predict(X_batch)
Y_predict[:10].data
Out[14]:
This example requires some learning of temporal dependencies. We want to learn our network the output when the input $x$'s bits are shifted right by $K$ positions. This can be done using x >> K. Similarly, a left shift can be done using (a << K) & (1 << N) -1), where $N$ is the max length of the bit sequence.
In [15]:
def create_dataset(max_len=5, K=4):
"""Create a dataset of max_len bits and their flipped values
Args:
max_len: Maximum number of bits in the number
"""
X, Y = [], []
max_val = 2**max_len
for i in range(max_val):
x = "{0:0{1}b}".format(i,max_len)
y = "{0:0{1}b}".format(i>>K,max_len)
x = tuple(map(int, x))
y = tuple(map(int, y))
X.append(x)
Y.append(y)
return X, Y
In [16]:
input_size=2
embedding_size=3
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
f = Input2Hidden(embedding_size, concat_layers=True)
g = Hidden2Output(embedding_size, 2, concat_layers=True)
rnn_cell = CustomRNNCell(f,g)
model = Model(embedding, rnn_cell)
In [17]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=128, shuffle=True)
In [18]:
%%time
learning_rate = 1e-4
max_epochs = 10000
check_every=500
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
loss = model.loss(*tensors2variables(X_tensors, Y_tensors))
Y_predict = model.predict(Variable(X_tensors)).data
accuracy = (Y_tensors == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
In [19]:
Y_predict = model.predict(X_batch)
Y_predict[:10].data
Out[19]:
The network trained above takes more than 10,000 epochs to converge to only $90\%$ accuracy. This reflects a major shortcoming of general RNN's. The shortcoming comes from a problem known as vanishing gradients, where gradients based on more distant steps become numerically too small to update the current layer, leading to information loss and failure to learn long range dependencies. Researcher's have worked around this using what is known as gated or memory based RNN cell's, which allows the information to be stored for a longer duration in the network and the gradients from long range dependencies to flow more easily. Two of the most popular variants are Long Short Term Memory (LSTM) cells and Gated Recurrent Unit (GRU) cells. The core idea is to allow some memory of the current state to be stored for the long time either in a seperate memory cell or in the hidden state. This is usually done by selectively reading and editing from the memory based on the current step. In the following sections we will understand the GRU cells which are a very simple extension of RNN and solve the vanishing gradient problem. The LSTM cells are a bit more involved and will be discussed later.
The idea behind GRU's is to update part of the hidden state and retain the rest. This is done using the following functions:
The implementation is as follow:
$$ \begin{equation} reset = \sigma(W_{r}[x_i, h_{i-1}])\\ update = \sigma(W_{u}[x_i, h_{i-1}])\\ interim\_hidden = tanh(W_{i}[x_i, reset \circ h_{i-1}])\\ h^{'}_{i} = update \circ interim\_hidden + (1-update) \circ h_{i-1} \\ \end{equation} $$
In [20]:
class GRUCell(nn.Module):
def __init__(self, input_dim, output_dim):
super(GRUCell, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.reset_linear = nn.Linear(2*self.input_dim, self.input_dim)
self.update_linear = nn.Linear(2*self.input_dim, self.input_dim)
self.interim_linear = nn.Linear(2*self.input_dim, self.input_dim)
self.output_linear = nn.Linear(2*self.input_dim, output_dim)
def forward(self, x, h):
concat_tensors = torch.cat([x,h], dim=1)
reset = F.sigmoid(self.reset_linear(concat_tensors))
update = F.sigmoid(self.update_linear(concat_tensors))
reset_hidden = reset * h
concat_reset_hidden = torch.cat([x, reset_hidden], dim=1)
interim_hidden = F.tanh(self.interim_linear(concat_reset_hidden))
h_prime = update * interim_hidden + (1-update) * h
concat_out = torch.cat([x, h_prime], dim=1)
output = F.tanh(self.output_linear(concat_out))
return output, h_prime
In [21]:
input_size=2
embedding_size=3
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
rnn_cell = GRUCell(embedding_size, output_size)
model = Model(embedding, rnn_cell)
In [22]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=64, shuffle=True)
In [23]:
%%time
learning_rate = 1e-4
max_epochs = 3000
check_every=500
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
loss = model.loss(*tensors2variables(X_tensors, Y_tensors))
Y_predict = model.predict(Variable(X_tensors)).data
accuracy = (Y_tensors == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
In [24]:
model
Out[24]:
In [25]:
input_size=2
embedding_size=3
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
rnn_cell = GRUCell(embedding_size, output_size)
model = Model(embedding, rnn_cell).cuda()
In [26]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
batch_size=64
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=batch_size, shuffle=True, pin_memory=True)
In [27]:
overall_hidden = Variable(torch.zeros(1,model.embedding.embedding_dim))
overall_hidden = overall_hidden.cuda()
overall_hidden
Out[27]:
In [28]:
%%time
learning_rate = 1e-4
max_epochs = 3000
check_every=500
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch = X_batch.cuda(async=True)
Y_batch = Y_batch.cuda(async=True)
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
hidden = overall_hidden.repeat(batch_size, 1)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch, hidden=hidden)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
hidden = overall_hidden.repeat(X_tensors.shape[0], 1)
loss = model.loss(*tensors2variables(X_tensors.cuda(), Y_tensors.cuda()),
hidden=hidden)
Y_predict = model.predict(Variable(X_tensors.cuda()), hidden=hidden).data
accuracy = (Y_tensors.cuda() == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
It looks like the GPU version is actually a bit slower in this case. This might be due to the small size of our dataset, for which the cost of moving tensors to GPU is greater than the gain by speeding up network computations.
In [29]:
input_size=2
embedding_size=10
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
rnn_cell = GRUCell(embedding_size, output_size)
model = Model(embedding, rnn_cell).cuda()
In [30]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
batch_size=64
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=batch_size, shuffle=True, pin_memory=True)
In [31]:
overall_hidden = Variable(torch.zeros(1,model.embedding.embedding_dim))
overall_hidden = overall_hidden.cuda()
overall_hidden
Out[31]:
In [32]:
%%time
learning_rate = 1e-4
max_epochs = 1000
check_every=50
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch = X_batch.cuda(async=True)
Y_batch = Y_batch.cuda(async=True)
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
hidden = overall_hidden.repeat(batch_size, 1)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch, hidden=hidden)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
hidden = overall_hidden.repeat(X_tensors.shape[0], 1)
loss = model.loss(*tensors2variables(X_tensors.cuda(), Y_tensors.cuda()),
hidden=hidden)
Y_predict = model.predict(Variable(X_tensors.cuda()), hidden=hidden).data
accuracy = (Y_tensors.cuda() == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
As we can see, the network converges 10x quicker than the one with lower capacity and also achieves $100\%$ accuracy in just 600 epochs. This is a very useful result, as it shows that in order to learn more complex functionalities we need networks with larger capacities as well as computationally efficient structures. Luckily for us many of the standard functionalities, are usually implemented efficiently in neural network libraries. Pytorch implements many of the standard neural network modules efficiently using it's C code, which can give us an order of magniture of improvement (especially for larger networks). These modules include GRU cells and a GRU module which can process the whole sequence. We will look at these in detail below.
In [33]:
class PytorchGRUCell(nn.Module):
def __init__(self, input_dim, output_dim):
super(PytorchGRUCell, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.gru = torch.nn.GRUCell(self.input_dim, self.input_dim)
self.output_linear = nn.Linear(2*self.input_dim, self.output_dim)
def forward(self, x, h):
h_prime = self.gru(x,h)
concat_out = torch.cat([x, h_prime], dim=1)
output = F.tanh(self.output_linear(concat_out))
return output, h_prime
In [34]:
input_size=2
embedding_size=10
output_size=2
embedding = nn.Embedding(input_size, embedding_size)
rnn_cell = PytorchGRUCell(embedding_size, output_size)
model = Model(embedding, rnn_cell).cuda()
In [35]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
batch_size=64
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=batch_size, shuffle=True, pin_memory=True)
In [36]:
overall_hidden = Variable(torch.zeros(1,model.embedding.embedding_dim))
overall_hidden = overall_hidden.cuda()
overall_hidden
Out[36]:
In [37]:
%%time
learning_rate = 1e-4
max_epochs = 1000
check_every=50
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch = X_batch.cuda(async=True)
Y_batch = Y_batch.cuda(async=True)
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
hidden = overall_hidden.repeat(batch_size, 1)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch, hidden=hidden)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
hidden = overall_hidden.repeat(X_tensors.shape[0], 1)
loss = model.loss(*tensors2variables(X_tensors.cuda(), Y_tensors.cuda()),
hidden=hidden)
Y_predict = model.predict(Variable(X_tensors.cuda()), hidden=hidden).data
accuracy = (Y_tensors.cuda() == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
Great, this implementation is almost 2x times faster than our implementation, probably because it is written using the C backend.
In [38]:
torch.cat([torch.zeros(1,2,3), torch.ones(1,2,3)], 2)
Out[38]:
In [39]:
class PyTorchModel(nn.Module):
def __init__(self, input_size, embedding_size, output_size):
super(PyTorchModel, self).__init__()
self.input_size = input_size
self.embedding_size = embedding_size
self.output_size = output_size
self.embedding = nn.Embedding(input_size, embedding_size)
self.rnn = nn.GRU(embedding_size, embedding_size)
self.output_linear = nn.Linear(2*self.embedding_size, self.output_size)
self.loss_function = nn.CrossEntropyLoss()
def forward(self, word_ids, hidden=None):
if hidden is None:
hidden = Variable(torch.zeros(
word_ids.data.shape[0],self.embedding.embedding_dim))
assert isinstance(hidden, Variable)
embeddings = self.embedding(word_ids)
max_seq_length = word_ids.data.shape[-1]
## RNN input and output shapes are (seq_len, batch_size, input_size)
embeddings = embeddings.permute(1,0,2)
hidden_states, hidden = self.rnn(embeddings)
concat_tensors = torch.cat([embeddings, hidden_states], 2)
concat_tensors = concat_tensors.permute(1,0,2).contiguous()
concat_tensors = concat_tensors.view(-1, concat_tensors.data.shape[2])
outputs = self.output_linear(concat_tensors)
hidden_states = hidden_states.permute(1,0,2)
outputs = outputs.view(self.input_size, -1, self.output_size)
assert isinstance(outputs, Variable)
assert isinstance(hidden_states, Variable)
return outputs, hidden_states
def loss(self, word_ids, target_ids, hidden=None):
outputs, hidden_states = self.forward(word_ids, hidden=hidden)
outputs = outputs.view(-1, outputs.data.shape[-1])
target_ids = target_ids.view(-1)
assert isinstance(outputs, Variable)
assert isinstance(target_ids, Variable)
#print("output={}\ttargets={}".format(outputs.data.shape,target_ids.data.shape))
loss = self.loss_function(outputs, target_ids)
return loss
def predict(self, word_ids, hidden=None):
outputs, hidden_states = self.forward(word_ids, hidden=hidden)
outputs = outputs.view(-1, outputs.data.shape[-1])
max_scores, predictions = outputs.max(1)
predictions = predictions.view(*word_ids.data.shape)
#print(word_ids.data.shape, predictions.data.shape)
assert word_ids.data.shape == predictions.data.shape, "word_ids: {}, predictions: {}".format(
word_ids.data.shape, predictions.data.shape
)
return predictions
In [40]:
input_size=2
embedding_size=10
output_size=2
model = PyTorchModel(input_size, embedding_size, output_size).cuda()
In [41]:
X, Y = create_dataset(max_len=10, K=4)
X_tensors, Y_tensors = tuple(map(torch.LongTensor, [X, Y]))
print(X_tensors.shape, Y_tensors.shape)
assert X_tensors.shape == Y_tensors.shape, "X and Y should be of same shape"
batch_size=64
train = data_utils.TensorDataset(X_tensors, Y_tensors)
train_loader = data_utils.DataLoader(train, batch_size=batch_size, shuffle=True, pin_memory=True)
In [42]:
overall_hidden = Variable(torch.zeros(1,model.embedding.embedding_dim))
overall_hidden = overall_hidden.cuda()
overall_hidden
Out[42]:
In [43]:
%%time
learning_rate = 1e-4
max_epochs = 1000
check_every=50
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(1, max_epochs+1):
for X_batch, Y_batch in train_loader:
X_batch = X_batch.cuda(async=True)
Y_batch = Y_batch.cuda(async=True)
X_batch, Y_batch = tensors2variables(X_batch, Y_batch)
hidden = overall_hidden.repeat(batch_size, 1)
# Forward pass: compute predicted y by passing x to the model.
loss = model.loss(X_batch, Y_batch, hidden=hidden)
# Before the backward pass, use the optimizer object to zero all of the
# gradients for the variables it will update (which are the learnable weights
# of the model)
optimizer.zero_grad()
# Backward pass: compute gradient of the loss with respect to model
# parameters
loss.backward()
# Calling the step function on an Optimizer makes an update to its
# parameters
optimizer.step()
if epoch % check_every != 0:
continue
hidden = overall_hidden.repeat(X_tensors.shape[0], 1)
loss = model.loss(*tensors2variables(X_tensors.cuda(), Y_tensors.cuda()),
hidden=hidden)
Y_predict = model.predict(Variable(X_tensors.cuda()), hidden=hidden).data
accuracy = (Y_tensors.cuda() == Y_predict).sum() /operator.mul(*Y_tensors.shape) * 100.
print("Epoch[{:03d}]: loss={:5.3f}; accuracy={:.3f}%".format(epoch, loss.data[0], accuracy))
This is super fast, 3x faster than using the GRUCell and 6x faster than our implementation.
This concludes our introduction to sequence tagging using Pytorch. The example covered here were very small so as to demonstrate the code required to implement a neural network as well as to give an intuition about the kind of tasks the networks can handle. More complex models can be built on top of this demo, which can handle variable length sequences, complex inference process (e.g. Linear Chain Conditional Random Fields for predicting the best sequence of outputs), and complex handling of input like words, phrases, etc.
In [44]:
print("Pytorch Version: {}".format(torch.__version__))
In [ ]: