64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
import torch
|
|
from aiRNN import dataloader, models, losses
|
|
import pathlib
|
|
|
|
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")
|
|
|
|
dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda")
|
|
results = []
|
|
|
|
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 r in results:
|
|
print(f"Model: {r[0]}, hidden_size={r[1]}, rnn_size={r[2]} => final loss={r[3]}") |