Add mode 4

Code/python/notebooks/hyperparameter_scan.py

@@ -158,6 +158,8 @@ def save_results(results, preliminary=False):
         cols = ["Model", "Hidden Layers", "RNN Layers", "Dropout", "Final Loss"]
     elif MODE == 3:
         cols = ["Model", "Warm-up Steps", "Prediction Steps", "Altitude Weight", "Final Loss"]
+    elif MODE == 4:
+        cols = ["Model", "Warm-up Steps", "Prediction Steps", "Final Loss"]
     else:
         raise ValueError("Invalid MODE")
 
@@ -314,6 +316,76 @@ def run_mode_3():
 
     return results
 
+# ------------------------------------------------------------
+# Mode 4 sweep
+# ------------------------------------------------------------
+def run_mode_4():
+    results = []
+
+    step = 10
+    base_name = "LSTM"
+    cls = models.ThreeInputLSTM
+    hidden_size = 16
+    rnn_size = 64
+    hidden_layers = 0
+    rnn_layers = 3
+    rnn_dropout = 0.0
+    altitude_weight = 1e-3
+    warm = 900
+    pred = 150
+    start_offset = 30 * 60 - warm
+    end_offset = 30 * 60 - pred
+
+    base_ds = dataloader.SaveDataset(
+            torch.load("long_dataset.pt"),
+            step=step,
+            start_offset=start_offset,
+            end_offset=end_offset,
+        )
+    train_ds, val_ds = split_dataset(base_ds)
+    model = cls(
+                    time_in=2,
+                    feat_in=4,
+                    context_in=5,
+                    hidden_size=hidden_size,
+                    rnn_size=rnn_size,
+                    out_size=3,
+                    hidden_layers=hidden_layers,
+                    rnn_layers=rnn_layers,
+                    rnn_dropout=rnn_dropout,
+                    device=DEVICE,
+                )
+    loss = train_one_model(model, train_ds, val_ds, warm // step, pred // step, altitude_weight=altitude_weight, batch_size=64, epochs=500)
+    print(f"Trained base model with loss {loss}")
+
+    for warm_up_test in range(300, 1800, 100):
+        for pred_steps_test in range(50, 1800, 100):
+            new_ds = dataloader.SaveDataset(
+                torch.load("long_dataset.pt"),
+                step=step,
+                start_offset=30 * 60 - warm_up_test,
+                end_offset=30 * 60 - pred_steps_test,
+            )
+            train_ds, val_ds = split_dataset(new_ds)
+            with torch.no_grad():
+                val_losses = []
+                val_loader = torch.utils.data.DataLoader(val_ds, batch_size=64, collate_fn=dataloader.collate_to_cpu, num_workers=4)
+                for X_f, X_t, y, X_c in val_loader:
+                    X_f = X_f.to(DEVICE)
+                    X_t = X_t.to(DEVICE)
+                    y = y.to(DEVICE)
+                    if X_c is not None:
+                        X_c = X_c.to(DEVICE)
+                    y_pred, _ = model(X_t, X_f, X_c, warm_up_test // step, pred_steps_test // step)
+                    val_losses.append(losses.HaversineAltitudeLoss(alt_const=altitude_weight)(y_pred, y).item())
+                val_loss = sum(val_losses) / len(val_losses)
+                print(f"Validation loss for warm_up={warm_up_test}, pred_steps={pred_steps_test}: {val_loss}")
+                results.append((base_name, warm_up_test, pred_steps_test, val_loss))
+                save_results(results, preliminary=True)
+
+
+
+
 
 # ------------------------------------------------------------
 # Dispatch
@@ -325,6 +397,8 @@ elif MODE == 2:
     results = run_mode_2()
 elif MODE == 3:
     results = run_mode_3()
+elif MODE == 4:
+    results = run_mode_4()
 else:
     raise ValueError("Invalid MODE.")