-
Notifications
You must be signed in to change notification settings - Fork 12
/
LSTMModel.py
23 lines (20 loc) · 1 KB
/
LSTMModel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import torch.nn as nn
class Lstm(nn.Module):
def __init__(self, input_size=8, hidden_size=32, num_layers=1 , output_size=1 , dropout=0, batch_first=True):
super(Lstm, self).__init__()
# Inputs #batch,seq_len, input_size
self.hidden_size = hidden_size
self.input_size = input_size
self.num_layers = num_layers
self.output_size = output_size
self.dropout = dropout
self.batch_first = batch_first
self.rnn = nn.LSTM(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=self.num_layers, batch_first=self.batch_first, dropout=self.dropout)
self.linear = nn.Linear(self.hidden_size, self.output_size)
def forward(self, x):
out, (hidden, cell) = self.rnn(x)
# x.shape : batch, seq_len, hidden_size , hn.shape and cn.shape : num_layes * direction_numbers, batch, hidden_size
a, b, c = hidden.shape
out = self.linear(hidden.reshape(a * b, c))
out = self.linear(hidden)
return out