File size: 2,056 Bytes
0b0933a |
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 |
import sys
import os
import json
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from models.forecasting.lstm import LSTMForecaster
df = pd.read_csv("../data/processed/merged_features.csv")
df = df.select_dtypes(include=[np.number]).dropna()
data = df.values
seq_len = 30
X, y = [], []
for i in range(len(data) - seq_len - 1):
X.append(data[i:i+seq_len])
y.append(data[i+seq_len][0])
X = torch.tensor(np.array(X), dtype=torch.float32)
y = torch.tensor(np.array(y), dtype=torch.float32).unsqueeze(1)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=32)
input_size = X.shape[2]
hidden_size = 256
num_layers = 2
output_size = 1
model = LSTMForecaster(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
output_size=output_size
)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
model.train()
total_loss = 0
for xb, yb in train_loader:
optimizer.zero_grad()
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch+1}: Train Loss = {avg_loss:.4f}")
os.makedirs("trained_models", exist_ok=True)
torch.save(model.state_dict(), "trained_models/lstm_forecaster.pt")
config = {
"input_size": input_size,
"hidden_size": hidden_size,
"num_layers": num_layers,
"output_size": output_size
}
with open("trained_models/config.json", "w") as f:
json.dump(config, f)
print("✅ Model trained and saved.")
|