Add some finess to the hyperparameter scan
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
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
|
||||
@@ -22,43 +23,52 @@ if not pathlib.Path("dataset.pt").exists():
|
||||
)
|
||||
dataset.save_entire_dataset("dataset.pt")
|
||||
|
||||
dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda")
|
||||
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)
|
||||
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 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]} => final loss={r[3]}")
|
||||
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)
|
||||
@@ -126,11 +126,12 @@ class BaseDataset(Dataset):
|
||||
torch.save(data_dict, filepath)
|
||||
|
||||
class SaveDataset(Dataset):
|
||||
def __init__(self, data_dict, device="cpu"):
|
||||
def __init__(self, data_dict, device="cpu", step=1):
|
||||
super().__init__()
|
||||
self.X_feat = data_dict["X_feat"].to(device)
|
||||
self.X_time = data_dict["X_time"].to(device)
|
||||
self.Y_out = data_dict["Y_out"].to(device)
|
||||
self.step = step
|
||||
self.device = device
|
||||
if "X_context" in data_dict:
|
||||
self.X_context = data_dict["X_context"].to(device)
|
||||
@@ -138,14 +139,14 @@ class SaveDataset(Dataset):
|
||||
self.X_context = None
|
||||
|
||||
def __len__(self):
|
||||
return self.X_feat.shape[0]
|
||||
return self.X_feat.shape[0] // self.step
|
||||
|
||||
def __getitem__(self, idx):
|
||||
X_f = self.X_feat[idx]
|
||||
X_t = self.X_time[idx]
|
||||
Y = self.Y_out[idx]
|
||||
X_f = self.X_feat[idx * self.step]
|
||||
X_t = self.X_time[idx * self.step]
|
||||
Y = self.Y_out[idx * self.step]
|
||||
if self.X_context is not None:
|
||||
X_c = self.X_context[idx]
|
||||
X_c = self.X_context[idx * self.step]
|
||||
else:
|
||||
X_c = None
|
||||
return X_f, X_t, Y, X_c
|
||||
|
||||
@@ -9,9 +9,10 @@ class BaseRNN(nn.Module):
|
||||
|
||||
base_unit = lambda in_size, out_size: nn.Sequential(
|
||||
nn.Linear(in_size, out_size, device=device),
|
||||
nn.Tanh(),
|
||||
nn.GELU(),
|
||||
nn.Linear(out_size, out_size, device=device),
|
||||
nn.Tanh(),
|
||||
nn.GELU(),
|
||||
nn.LayerNorm(out_size, device=device)
|
||||
)
|
||||
|
||||
self.time_proj = base_unit(time_in, hidden_size)
|
||||
@@ -48,11 +49,7 @@ class BaseRNN(nn.Module):
|
||||
raise ValueError("Unsupported rnn_type")
|
||||
|
||||
readout_in = rnn_size
|
||||
self.readout = nn.Sequential(
|
||||
nn.Linear(readout_in, hidden_size, device=device),
|
||||
nn.Tanh(),
|
||||
nn.Linear(hidden_size, out_size, device=device),
|
||||
)
|
||||
self.readout = nn.Linear(readout_in, out_size, device=device)
|
||||
self.context_in = context_in
|
||||
self.rnn_size = rnn_size
|
||||
|
||||
|
||||
Reference in New Issue
Block a user