Add some finess to the hyperparameter scan

Code/python/notebooks/hyperparameter_scan.py

@@ -1,6 +1,7 @@
 import torch
 from aiRNN import dataloader, models, losses
 import pathlib
+import pandas as pd
 
 file_list = pathlib.Path("../../cpp/known_routes_and_aircraft.csv")
 base_path = file_list.parent
@@ -22,43 +23,52 @@ if not pathlib.Path("dataset.pt").exists():
   )
   dataset.save_entire_dataset("dataset.pt")
 
-dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda")
 results = []
+for step in [1, 5, 10, 30]:
+    dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda", step=step)
 
-for base_name, base_model in [
-    ("RNN", models.ThreeInputRNN),
-    ("LSTM", models.ThreeInputLSTM),
-    ("GRU", models.ThreeInputGRU),
-]:
-    for hidden_size in [16, 32, 64]:
-        for rnn_size in [32, 64, 128]:
-            model = base_model(
-                time_in=2,
-                feat_in=4,
-                context_in=5,
-                hidden_size=hidden_size,
-                rnn_size=rnn_size,
-                out_size=3,
-                device="cuda",
-            )
-            print(f"Training {base_name} with hidden_size={hidden_size}, rnn_size={rnn_size}")
-            optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
-            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
-            criterion = losses.HaversineAltitudeLoss(alt_const=1e-3)
-            for epoch in range(100):
-                loss_history = []
-                for X_f, X_t, y, X_c in torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True):
-                    optimizer.zero_grad()
-                    y_pred, _ = model(X_t, X_f, X_c, 450, 150)
-                    loss = criterion(y_pred, y)
-                    loss.backward()
-                    optimizer.step()
-                    loss_history.append(loss.item())
-                loss = sum(loss_history) / len(loss_history)
-                scheduler.step(loss)
-                print(f"Epoch {epoch}: loss={loss}, lr={optimizer.param_groups[0]['lr']}")
-            torch.save(model.state_dict(), f"{base_name}_hs{hidden_size}_rs{rnn_size}.pt")
-            results.append((base_name, hidden_size, rnn_size, loss))
+    for base_name, base_model in [
+        ("RNN", models.ThreeInputRNN),
+        ("LSTM", models.ThreeInputLSTM),
+        ("GRU", models.ThreeInputGRU),
+    ]:
+        for hidden_size in [16, 32, 64]:
+            for rnn_size in [32, 64, 128]:
+                model = base_model(
+                    time_in=2,
+                    feat_in=4,
+                    context_in=5,
+                    hidden_size=hidden_size,
+                    rnn_size=rnn_size,
+                    out_size=3,
+                    device="cuda",
+                )
+                print(f"Training {base_name} with hidden_size={hidden_size}, rnn_size={rnn_size}")
+                optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
+                scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
+                criterion = losses.HaversineAltitudeLoss(alt_const=1e-3)
+                loss = -1.0
+                for epoch in range(100):
+                    loss_history = []
+                    for X_f, X_t, y, X_c in torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True):
+                        optimizer.zero_grad()
+                        warm_up_steps = 450 // step
+                        pred_steps = 150 // step
+                        y_pred, _ = model(X_t, X_f, X_c, warm_up_steps, pred_steps)
+                        loss = criterion(y_pred, y)
+                        loss.backward()
+                        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                        optimizer.step()
+                        loss_history.append(loss.item())
+                    loss = sum(loss_history) / len(loss_history)
+                    scheduler.step(loss)
+                    print(f"Epoch {epoch}: loss={loss}, lr={optimizer.param_groups[0]['lr']}")
+                torch.save(model.state_dict(), f"{base_name}_hs{hidden_size}_rs{rnn_size}_step{step}.pt")
+                results.append((base_name, hidden_size, rnn_size, step, loss))
 
 for r in results:
-    print(f"Model: {r[0]}, hidden_size={r[1]}, rnn_size={r[2]} => final loss={r[3]}")
+    print(f"Model: {r[0]}, hidden_size={r[1]}, rnn_size={r[2]}, step={r[3]} => final loss={r[4]}")
+
+# Save results to CSV
+df = pd.DataFrame(results, columns=["Model", "Hidden Size", "RNN Size", "Step", "Final Loss"])
+df.to_csv("hyperparameter_scan_results.csv", index=False)