lars/aiRtrafficNN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add more hyperparameter scan modes

Browse Source
This commit is contained in:
Lars Bogner
2025-11-25 09:38:25 +01:00
parent b012a86768
commit c0b2545671
5 changed files with 1152 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Code/python/notebooks/hyperparameter_scan.py
+268 -73
View File
@@ -1,90 +1,285 @@
import torch
from aiRNN import dataloader, models, losses
import pathlib
import pandas as pd
from aiRNN import dataloader, models, losses
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()]
MODE = 1
DEVICE = "cuda"
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")
results = []
for step in [1, 5, 10, 30]:
dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda", step=step)
dataset_length = len(dataset)
dataset, val_dataset = torch.utils.data.random_split(
# ------------------------------------------------------------
# Dataset creation
# ------------------------------------------------------------
def load_file_list(path: pathlib.Path):
base = path.parent
return [
(base / f).resolve()
for f in path.read_text().splitlines()
if (base / f).exists()
]
def ensure_dataset(path, **kwargs):
if not pathlib.Path(path).exists():
ds = dataloader.EvenlySpacedDataset(**kwargs, device=DEVICE)
ds.save_entire_dataset(path)
file_list_path = pathlib.Path("../../cpp/known_routes_and_aircraft.csv")
file_list = load_file_list(file_list_path)
ensure_dataset(
"dataset.pt",
filepaths=file_list,
n_input=30 * 15,
n_output=30 * 5,
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"),
)
ensure_dataset(
"long_dataset.pt",
filepaths=file_list,
n_input=30 * 60,
n_output=30 * 60,
n_windows_per_file=5,
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"),
)
# ------------------------------------------------------------
# Training utilities
# ------------------------------------------------------------
def split_dataset(dataset, frac=0.8, seed=42):
n = len(dataset)
train_n = int(frac * n)
return torch.utils.data.random_split(
dataset,
[int(0.8 * dataset_length), dataset_length - int(0.8 * dataset_length)],
generator=torch.Generator().manual_seed(42)
[train_n, n - train_n],
generator=torch.Generator().manual_seed(seed),
)
print(f"Starting hyperparameter scan for 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(
def train_one_model(model, dataset, val_dataset, warm_up_steps, pred_steps, epochs=100, altitude_weight=1e-3):
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=altitude_weight)
for epoch in range(epochs):
batch_losses = []
loader = torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True, collate_fn=dataloader.collate_to_cuda(DEVICE), num_workers=4)
for X_f, X_t, y, X_c in loader:
optimizer.zero_grad()
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()
batch_losses.append(loss.item())
mean_loss = sum(batch_losses) / len(batch_losses)
scheduler.step(mean_loss)
print(f"Epoch {epoch}: loss={mean_loss}, lr={optimizer.param_groups[0]['lr']}")
# Validation
with torch.no_grad():
val_losses = []
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=256, collate_fn=dataloader.collate_to_cuda(DEVICE), num_workers=4)
for X_f, X_t, y, X_c in val_loader:
y_pred, _ = model(X_t, X_f, X_c, warm_up_steps, pred_steps)
val_losses.append(criterion(y_pred, y).item())
val_loss = sum(val_losses) / len(val_losses)
print(f"Validation loss: {val_loss}")
return val_loss
def save_results(results, preliminary=False):
if MODE == 1:
cols = ["Model", "Hidden Size", "RNN Size", "Step", "Final Loss"]
elif MODE == 2:
cols = ["Model", "Hidden Layers", "RNN Layers", "Dropout", "Final Loss"]
elif MODE == 3:
cols = ["Model", "Warm-up Steps", "Prediction Steps", "Altitude Weight", "Final Loss"]
else:
raise ValueError("Invalid MODE")
df = pd.DataFrame(results, columns=cols)
suffix = "_preliminary" if preliminary else ""
df.to_csv(f"hyperparameter_scan_results_mode{MODE}{suffix}.csv", index=False)
# ------------------------------------------------------------
# MODE 1 sweep
# ------------------------------------------------------------
def run_mode_1():
results = []
for step in [1, 5, 10, 30]:
base_ds = dataloader.SaveDataset(torch.load("dataset.pt"), step=step)
train_ds, val_ds = split_dataset(base_ds)
print(f"Starting hyperparameter scan for step={step}")
for name, cls 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 = cls(
time_in=2,
feat_in=4,
context_in=5,
hidden_size=hidden_size,
rnn_size=rnn_size,
out_size=3,
device=DEVICE,
)
warm = 450 // step
pred = 150 // step
print(f"Training {name} hs={hidden_size} rs={rnn_size}")
loss = train_one_model(model, train_ds, val_ds, warm, pred)
torch.save(model.state_dict(), f"{name}_hs{hidden_size}_rs{rnn_size}_step{step}.pt")
results.append((name, hidden_size, rnn_size, step, loss))
save_results(results, preliminary=True)
return results
# ------------------------------------------------------------
# MODE 2 sweep
# ------------------------------------------------------------
def run_mode_2():
results = []
step = 1
base_ds = dataloader.SaveDataset(torch.load("dataset.pt"), step=step)
train_ds, val_ds = split_dataset(base_ds)
base_name = "GRU"
cls = models.ThreeInputGRU
hidden_size = 32
rnn_size = 64
for hl in [0, 1, 2, 4]:
for rl in [2, 3, 4, 5]:
for dropout in [0.0, 0.1, 0.2, 0.3]:
model = cls(
time_in=2,
feat_in=4,
context_in=5,
hidden_size=hidden_size,
rnn_size=rnn_size,
out_size=3,
device="cuda",
hidden_layers=hl,
rnn_layers=rl,
rnn_dropout=dropout,
device=DEVICE,
)
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']}")
with torch.no_grad():
val_loss_history = []
for X_f, X_t, y, X_c in torch.utils.data.DataLoader(val_dataset, batch_size=256):
warm_up_steps = 450 // step
pred_steps = 150 // step
y_pred, _ = model(X_t, X_f, X_c, warm_up_steps, pred_steps)
v_loss = criterion(y_pred, y)
val_loss_history.append(v_loss.item())
loss = sum(val_loss_history) / len(val_loss_history)
print(f"Validation loss: {loss}")
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]}, step={r[3]} => final loss={r[4]}")
warm = 450 // step
pred = 150 // step
# 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)
print(f"Training {base_name} hl={hl} rl={rl} do={dropout}")
loss = train_one_model(model, train_ds, val_ds, warm, pred)
torch.save(model.state_dict(), f"{base_name}_hl{hl}_rl{rl}_do{int(dropout*10)}.pt")
results.append((base_name, hl, rl, dropout, loss))
save_results(results, preliminary=True)
return results
# ------------------------------------------------------------
# MODE 3 sweep
# ------------------------------------------------------------
def run_mode_3():
results = []
step = 1
base_name = "GRU"
cls = models.ThreeInputGRU
hidden_size = 32
rnn_size = 64
hidden_layers = 1
rnn_layers = 3
rnn_dropout = 0.1
for warm in [300, 600, 900, 1800]:
for pred in [150, 300, 600, 900, 1800]:
start = 30 * 60 - warm
end = 30 * 60 - pred
base_ds = dataloader.SaveDataset(
torch.load("long_dataset.pt"),
step=step,
start_offset=start,
end_offset=end,
)
train_ds, val_ds = split_dataset(base_ds)
for altitude_weight in [1e-5, 1e-4, 1e-3]:
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,
)
print(f"Training {base_name} warm={warm} pred={pred} altitude_weight={altitude_weight}")
loss = train_one_model(model, train_ds, val_ds, warm // step, pred // step, altitude_weight=altitude_weight)
torch.save(model.state_dict(), f"{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.pt")
results.append((base_name, warm, pred, altitude_weight, loss))
save_results(results, preliminary=True)
return results
# ------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------
if MODE == 1:
results = run_mode_1()
elif MODE == 2:
results = run_mode_2()
elif MODE == 3:
results = run_mode_3()
else:
raise ValueError("Invalid MODE.")
save_results(results)