File size: 1,997 Bytes
4ca4e4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | import torch
import math
import torch.nn as nn
class LSTM(nn.Module):
def __init__(self, embedding_dim, vocab_size, num_layers, dropout_rate=0.65):
super(LSTM, self).__init__()
self.num_layers = num_layers
self.n_embd = embedding_dim
self.word_embeddings = nn.Embedding(vocab_size, self.n_embd)
# The LSTM takes word embeddings as inputs, and outputs hidden states
# with dimensionality hidden_dim.
self.lstm = nn.LSTM(input_size=self.n_embd,
hidden_size=self.n_embd,
num_layers=self.num_layers,
batch_first=True,
dropout=dropout_rate)
# The linear layer that maps from hidden state space to tag space
self.head = nn.Linear(self.n_embd, vocab_size)
self.init_weights()
def forward(self, sentence, prev_state):
embeds = self.word_embeddings(sentence)
lstm_out, state = self.lstm(embeds, prev_state)
logits = self.head(lstm_out)
return logits, state
def init_weights(self):
init_range_emb = 0.1
init_range_other = 1/math.sqrt(self.n_embd)
self.word_embeddings.weight.data.uniform_(-init_range_emb, init_range_emb)
self.head.weight.data.uniform_(-init_range_other, init_range_other)
self.head.bias.data.zero_()
for i in range(self.num_layers):
self.lstm.all_weights[i][0] = torch.FloatTensor(self.n_embd,
self.n_embd).uniform_(-init_range_other, init_range_other)
self.lstm.all_weights[i][1] = torch.FloatTensor(self.n_embd,
self.n_embd).uniform_(-init_range_other, init_range_other)
def init_hidden(self, batch_size, device):
hidden = torch.zeros(self.num_layers, batch_size, self.n_embd).to(device)
cell = torch.zeros(self.num_layers, batch_size, self.n_embd).to(device)
return hidden, cell
|