Added LSTM model

This commit is contained in:
Victor Mylle
2023-11-28 22:27:15 +00:00
parent ffa19592f9
commit eba10c8f83
6 changed files with 117 additions and 38 deletions

45
src/models/lstm_model.py Normal file
View File

@@ -0,0 +1,45 @@
import torch
class LSTMModel(torch.nn.Module):
def __init__(self, inputSize, output_size, num_layers: int, hidden_size: int, dropout: float = 0.2):
super(LSTMModel, self).__init__()
self.inputSize = inputSize
self.output_size = output_size
self.num_layers = num_layers
self.hidden_size = hidden_size
self.dropout = dropout
self.lstm = torch.nn.LSTM(input_size=inputSize[-1], hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, batch_first=True)
self.linear = torch.nn.Linear(hidden_size, output_size)
def forward(self, x):
# Forward pass through the LSTM layers
_, (hidden_state, _) = self.lstm(x)
# Use the hidden state from the last time step for the output
output = self.linear(hidden_state[-1])
return output
class GRUModel(torch.nn.Module):
def __init__(self, inputSize, output_size, num_layers: int, hidden_size: int, dropout: float = 0.2):
super(GRUModel, self).__init__()
self.inputSize = inputSize
self.output_size = output_size
self.num_layers = num_layers
self.hidden_size = hidden_size
self.dropout = dropout
self.gru = torch.nn.GRU(input_size=inputSize[-1], hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, batch_first=True)
self.linear = torch.nn.Linear(hidden_size, output_size)
def forward(self, x):
# Forward pass through the GRU layers
_, hidden_state = self.gru(x)
# Use the hidden state from the last time step for the output
output = self.linear(hidden_state[-1])
return output