File for hyperparameter_scan

2025-11-23 12:03:03 +01:00
parent 7960c05bef
commit b47bc1e297
1 changed files with 64 additions and 0 deletions
@@ -0,0 +1,64 @@
+import torch
+from aiRNN import dataloader, models, losses
+import pathlib
+
+file_list = pathlib.Path("../../cpp/known_routes_and_aircraft.csv")
+base_path = file_list.parent
+file_list = file_list.read_text().splitlines()
+file_list = [(base_path / f).resolve() for f in file_list if (base_path / f).exists()]
+
+if not pathlib.Path("dataset.pt").exists():
+  dataset = dataloader.EvenlySpacedDataset(
+      filepaths=file_list,
+      n_input=30*15, # 15 minutes input
+      n_output=30*5,  # 5 minutes output
+      n_windows_per_file=7,
+      step=1,
+      feature_columns=("lat", "lon", "alt", "ias"),
+      context_columns=("last_lat", "last_lon", "last_alt", "last_ias", "last_timestamp"),
+      time_columns=("timestamp", "dt"),
+      target_columns=("lat", "lon", "alt"),
+      device="cuda",
+  )
+  dataset.save_entire_dataset("dataset.pt")
+
+dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda")
+results = []
+
+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 r in results:
+    print(f"Model: {r[0]}, hidden_size={r[1]}, rnn_size={r[2]} => final loss={r[3]}")