Add some finess to the hyperparameter scan

This commit is contained in:
2025-11-23 12:45:39 +01:00
parent b47bc1e297
commit a77fd08c5d
3 changed files with 57 additions and 49 deletions
+46 -36
View File
@@ -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)
+7 -6
View File
@@ -126,11 +126,12 @@ class BaseDataset(Dataset):
torch.save(data_dict, filepath)
class SaveDataset(Dataset):
def __init__(self, data_dict, device="cpu"):
def __init__(self, data_dict, device="cpu", step=1):
super().__init__()
self.X_feat = data_dict["X_feat"].to(device)
self.X_time = data_dict["X_time"].to(device)
self.Y_out = data_dict["Y_out"].to(device)
self.step = step
self.device = device
if "X_context" in data_dict:
self.X_context = data_dict["X_context"].to(device)
@@ -138,14 +139,14 @@ class SaveDataset(Dataset):
self.X_context = None
def __len__(self):
return self.X_feat.shape[0]
return self.X_feat.shape[0] // self.step
def __getitem__(self, idx):
X_f = self.X_feat[idx]
X_t = self.X_time[idx]
Y = self.Y_out[idx]
X_f = self.X_feat[idx * self.step]
X_t = self.X_time[idx * self.step]
Y = self.Y_out[idx * self.step]
if self.X_context is not None:
X_c = self.X_context[idx]
X_c = self.X_context[idx * self.step]
else:
X_c = None
return X_f, X_t, Y, X_c
+4 -7
View File
@@ -9,9 +9,10 @@ class BaseRNN(nn.Module):
base_unit = lambda in_size, out_size: nn.Sequential(
nn.Linear(in_size, out_size, device=device),
nn.Tanh(),
nn.GELU(),
nn.Linear(out_size, out_size, device=device),
nn.Tanh(),
nn.GELU(),
nn.LayerNorm(out_size, device=device)
)
self.time_proj = base_unit(time_in, hidden_size)
@@ -48,11 +49,7 @@ class BaseRNN(nn.Module):
raise ValueError("Unsupported rnn_type")
readout_in = rnn_size
self.readout = nn.Sequential(
nn.Linear(readout_in, hidden_size, device=device),
nn.Tanh(),
nn.Linear(hidden_size, out_size, device=device),
)
self.readout = nn.Linear(readout_in, out_size, device=device)
self.context_in = context_in
self.rnn_size = rnn_size