Add new multilayer model

Changes

1 changed files (+32/-0)

• model.py (new)
  • @@ -0,0 +1,32 @@
    + import torch.nn as nn + + + class Model(nn.Module): + def __init__(self, n_vocab, seq_size, embedding_size, lstm_size, num_layers, dropout): + super(Model, self).__init__() + + self.seq_size = seq_size + self.lstm_size = lstm_size + self.num_layers = num_layers + + self.embedding = nn.Embedding(n_vocab, embedding_size) + + self.lstm = nn.LSTM( + input_size=embedding_size, + hidden_size=lstm_size, + num_layers=num_layers, + batch_first=True, + dropout=dropout + ) + self.dense = nn.Linear(lstm_size, n_vocab) + + def forward(self, x, prev_state): + embed = self.embedding(x) + output, state = self.lstm(embed, prev_state) + logits = self.dense(output) + + return logits, state + + def zero_state(self, batch_size): + return (torch.zeros(self.num_layers, batch_size, self.lstm_size), + torch.zeros(self.num_layers, batch_size, self.lstm_size))