Update core/train_eval.py

core/train_eval.py

@@ -9,11 +9,11 @@ import matplotlib.pyplot as plt
 import os
 
 
-def create_sequences(data, window_size):
+def create_sequences(data, window_size, horizon=1):
     X, y = [], []
-    for i in range(len(data) - window_size):
+    for i in range(len(data) - window_size - horizon + 1):
         X.append(data[i:i + window_size])
-        y.append(data[i + window_size])
+        y.append(data[i + window_size:i + window_size + horizon].flatten())
     return np.array(X), np.array(y)
 
 
@@ -38,30 +38,41 @@ def train_and_evaluate(
     # Step 2: Normalize data
     scaler = StandardScaler()
     scaled_data = scaler.fit_transform(original_values.reshape(-1, 1))
-    X, y = create_sequences(scaled_data, window)
+    X, y = create_sequences(scaled_data, window, horizon)
 
-    # Step 3: Split
+    # Step 3: Split into train, validation, and test
     split = int(len(X) * (1 - test_split))
-    X_train, X_test = X[:split], X[split:]
-    y_train, y_test = y[:split], y[split:]
+    val_split = int(split * 0.9)  # 90% of training data for training, 10% for validation
+    X_train, X_val, X_test = X[:val_split], X[val_split:split], X[split:]
+    y_train, y_val, y_test = y[:val_split], y[val_split:split], y[split:]
 
     X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
     y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
+    X_val_tensor = torch.tensor(X_val, dtype=torch.float32)
+    y_val_tensor = torch.tensor(y_val, dtype=torch.float32)
     X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
     y_test_tensor = torch.tensor(y_test, dtype=torch.float32)
 
     train_loader = DataLoader(TensorDataset(X_train_tensor, y_train_tensor), batch_size=32, shuffle=True)
+    val_loader = DataLoader(TensorDataset(X_val_tensor, y_val_tensor), batch_size=32, shuffle=False)
     test_loader = DataLoader(TensorDataset(X_test_tensor, y_test_tensor), batch_size=32, shuffle=False)
 
     # Step 4: Model
     input_dim = X_train.shape[2] if X_train.ndim == 3 else 1
     model = model_cls(input_size=input_dim, hidden_size=hidden, num_layers=layers, output_size=horizon).to(device)
-    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0.01)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=0.01)
     loss_fn = nn.MSELoss()
 
     train_losses = []
+    val_losses = []
+    best_val_loss = float('inf')
+    patience = 5
+    counter = 0
+    best_model_state = None
+
     model.train()
     for epoch in range(epochs):
+        # Training
         epoch_loss = 0.0
         for xb, yb in train_loader:
             xb, yb = xb.to(device), yb.to(device)
@@ -72,10 +83,38 @@ def train_and_evaluate(
             optimizer.step()
             epoch_loss += loss.item()
         train_losses.append(epoch_loss / len(train_loader))
+
+        # Validation
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for xb, yb in val_loader:
+                xb, yb = xb.to(device), yb.to(device)
+                out = model(xb)
+                loss = loss_fn(out, yb)
+                val_loss += loss.item()
+        val_loss /= len(val_loader)
+        val_losses.append(val_loss)
+
         if verbose and (epoch + 1) % 10 == 0:
-            print(f"Epoch {epoch+1}/{epochs} - Loss: {train_losses[-1]:.4f}")
+            print(f"Epoch {epoch+1}/{epochs} - Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}")
+
+        # Early stopping
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            counter = 0
+            best_model_state = model.state_dict()
+        else:
+            counter += 1
+            if counter >= patience:
+                print(f"Early stopping at epoch {epoch+1}")
+                break
+
+    if best_model_state:
+        model.load_state_dict(best_model_state)
 
     result["train_loss"] = train_losses
+    result["val_loss"] = val_losses
 
     # Step 5: Evaluate
     model.eval()
@@ -90,16 +129,14 @@ def train_and_evaluate(
     preds = np.concatenate(preds, axis=0)
     targets = np.concatenate(targets, axis=0)
 
-    preds_inv = scaler.inverse_transform(preds.reshape(-1, 1)).reshape(preds.shape)
-    targets_inv = scaler.inverse_transform(targets.reshape(-1, 1)).reshape(targets.shape)
+    preds_inv = scaler.inverse_transform(preds.reshape(-1, horizon)).reshape(preds.shape)
+    targets_inv = scaler.inverse_transform(targets.reshape(-1, horizon)).reshape(targets.shape)
 
     mse = mean_squared_error(targets_inv, preds_inv)
     rmse = np.sqrt(mse)
     mae = mean_absolute_error(targets_inv, preds_inv)
-    # Calculate R2 score
     r2 = r2_score(targets_inv, preds_inv)
 
-
     result["metrics"] = {
         #"R2": round(r2, 4),
         "RMSE": round(rmse, 4),
@@ -110,17 +147,14 @@ def train_and_evaluate(
     result["actual"] = targets_inv
     result["predicted"] = result["forecast"]
 
-    # Step 6: Predict the next value (for UI)
+    # Step 6: Predict the next value(s) (for UI)
     latest_window = scaled_data[-window:].reshape(1, window, 1)
     latest_input = torch.tensor(latest_window, dtype=torch.float32).to(device)
 
     with torch.no_grad():
         future_pred = model(latest_input).cpu().numpy()
-    future_pred_inv = scaler.inverse_transform(future_pred.reshape(-1, 1)).reshape(future_pred.shape)
+    future_pred_inv = scaler.inverse_transform(future_pred.reshape(-1, horizon)).reshape(future_pred.shape)
 
-    result["latest_prediction"] = {
-        "input_index": len(original_values),
-        "value": float(future_pred_inv[0][0])
-    }
+    result["latest_prediction"] = future_pred_inv[0].tolist()  # List of horizon predictions
 
     return result