Lars Bogner
296 lines
9.2 KiB
Python
296 lines
9.2 KiB
Python
import torch
|
|
import pathlib
|
|
import pandas as pd
|
|
from aiRNN import dataloader, models, losses
|
|
|
|
MODE = 2
|
|
DEVICE = "cuda:2"
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# 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, batch_size=256):
|
|
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=batch_size, 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=batch_size, 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 = 10
|
|
base_ds = dataloader.SaveDataset(torch.load("dataset.pt"), step=step)
|
|
train_ds, val_ds = split_dataset(base_ds)
|
|
|
|
base_name = "LSTM"
|
|
cls = models.ThreeInputLSTM
|
|
hidden_size = 16
|
|
rnn_size = 64
|
|
|
|
for hl in [0, 1, 2, 4]:
|
|
for rl in [2, 3, 4, 5]:
|
|
for dropout in [0.0, 0.05, 0.1, 0.15]:
|
|
|
|
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 = 10
|
|
base_name = "LSTM"
|
|
cls = models.ThreeInputLSTM
|
|
hidden_size = 16
|
|
rnn_size = 64
|
|
hidden_layers = 1
|
|
rnn_layers = 3
|
|
rnn_dropout = 0.0
|
|
|
|
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, batch_size=64)
|
|
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)
|