-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
89 lines (61 loc) · 3.01 KB
/
models.py
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import torch
from torch import nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #selecting device
class RNNModel(nn.Module):
def __init__(self,input_size, hidden_size, rnn_layers, fc_hidden_dim):
super(RNNModel, self).__init__()
self.rnn_layers = rnn_layers
self.hidden_size = hidden_size
self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=rnn_layers, batch_first=True)
self.fcl = nn.Sequential(
nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),
nn.ReLU(),
nn.Linear(in_features=fc_hidden_dim, out_features=input_size)
)
def forward(self, x, prev_state):
out, state = self.rnn(x, prev_state)
fcl_out = self.fcl(out.reshape(-1,out.size(2)))
return fcl_out, state
def init_hidden(self, batch_size):
return torch.zeros((self.rnn_layers, batch_size, self.hidden_size), device=device)
def getModelName(self):
return "RNN"
class LSTMModel(nn.Module):
def __init__(self,input_size, hidden_size, lstm_layers, fc_hidden_dim):
super(LSTMModel, self).__init__()
self.lstm_layers = lstm_layers
self.hidden_size = hidden_size
self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=lstm_layers, batch_first=True)
self.fcl = nn.Sequential(
nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),
nn.ReLU(),
nn.Linear(in_features=fc_hidden_dim, out_features=input_size)
)
def forward(self, x, prev_state):
out, state = self.lstm(x, prev_state)
fcl_out = self.fcl(out.reshape(-1,out.size(2)))
return fcl_out, state
def init_hidden(self, batch_size):
return [torch.zeros((self.lstm_layers, batch_size, self.hidden_size), device=device),
torch.zeros((self.lstm_layers, batch_size, self.hidden_size), device=device)]
def getModelName(self):
return "LSTM"
class GRUModel(nn.Module):
def __init__(self,input_size, hidden_size, gru_layers, fc_hidden_dim):
super(GRUModel, self).__init__()
self.gru_layers = gru_layers
self.hidden_size = hidden_size
self.gru = nn.GRU(input_size=input_size, hidden_size=hidden_size, num_layers=gru_layers, batch_first=True)
self.fcl = nn.Sequential(
nn.Linear(in_features=hidden_size, out_features=fc_hidden_dim),
nn.ReLU(),
nn.Linear(in_features=fc_hidden_dim, out_features=input_size)
)
def forward(self, x, prev_state):
out, state = self.gru(x, prev_state)
fcl_out = self.fcl(out.reshape(-1,out.size(2)))
return fcl_out, state
def init_hidden(self, batch_size):
return torch.zeros((self.gru_layers, batch_size, self.hidden_size), device=device)
def getModelName(self):
return "GRU"