74 lines
3.2 KiB
Python
74 lines
3.2 KiB
Python
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
|
|
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")
|
|
|
|
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)
|
|
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]}, 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) |