|
import torch |
|
import torch.nn as nn |
|
|
|
class RNN_GEN(nn.Module): |
|
def __init__(self, num_countries, input_dim, output_dim, hidden_dim=128): |
|
super().__init__() |
|
self.input_to_hidden = nn.Linear(num_countries + input_dim + hidden_dim, hidden_dim) |
|
self.input_to_output = nn.Linear(num_countries + input_dim + hidden_dim, output_dim) |
|
self.output_to_output = nn.Linear(hidden_dim + output_dim, output_dim) |
|
self.log_softmax = nn.LogSoftmax(dim=1) |
|
self.hidden_dim = hidden_dim |
|
|
|
def forward(self, inp, country_tensor, hidden): |
|
combo = torch.cat([country_tensor, inp, hidden], dim=1) |
|
hidden = self.input_to_hidden(combo) |
|
out = self.input_to_output(combo) |
|
out_combo = torch.cat([hidden, out], dim=1) |
|
out = self.output_to_output(out_combo) |
|
out = self.log_softmax(out) |
|
return out, hidden |
|
|
|
def init_hidden(self): |
|
return torch.zeros((1, self.hidden_dim)) |
|
|
|
|
|
|