File size: 705 Bytes
7d950d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import torch
import torch.nn as nn
class LSTMPredictor(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, forecast_horizon=3, n_layers=2, dropout=0.2):
super().__init__()
self.hidden_dim = hidden_dim
self.n_layers = n_layers
self.forecast_horizon = forecast_horizon
self.lstm = nn.LSTM(input_dim, hidden_dim, n_layers, batch_first=True, dropout=dropout)
self.fc = nn.Linear(hidden_dim, output_dim * forecast_horizon)
def forward(self, x):
lstm_out, _ = self.lstm(x)
predictions = self.fc(lstm_out[:, -1, :])
return predictions.view(-1, self.forecast_horizon, predictions.shape[-1] // self.forecast_horizon)
|