import torch import pathlib import pandas as pd from aiRNN import dataloader, models, losses MODE = 1 DEVICE = "cuda" # ------------------------------------------------------------ # 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, [train_n, n - train_n], generator=torch.Generator().manual_seed(seed), ) 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_cpu, num_workers=4) for X_f, X_t, y, X_c in 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) 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_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_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, hidden_layers=hl, rnn_layers=rl, rnn_dropout=dropout, device=DEVICE, ) warm = 450 // step pred = 150 // step 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)