79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
import torch
|
|
import pathlib
|
|
import pandas as pd
|
|
from aiRNN import dataloader, models, losses
|
|
import numpy as np
|
|
import copy
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
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),
|
|
)
|
|
|
|
step = 10
|
|
base_name = "LSTM"
|
|
cls = models.ThreeInputLSTM
|
|
hidden_size = 16
|
|
rnn_size = 64
|
|
hidden_layers = 0
|
|
rnn_layers = 3
|
|
rnn_dropout = 0.0
|
|
altitude_weight = 1e-3
|
|
warm = 900
|
|
pred = 1600
|
|
start_offset = 30 * 60 - warm
|
|
end_offset = 30 * 60 - pred
|
|
base_ds = dataloader.SaveDataset(
|
|
torch.load("long_dataset.pt"),
|
|
step=step,
|
|
start_offset=start_offset,
|
|
end_offset=end_offset,
|
|
)
|
|
train_ds, val_ds = split_dataset(base_ds)
|
|
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,
|
|
)
|
|
|
|
model.load_state_dict(torch.load("LSTM_wu900_ps150_aw0.001.pt"))
|
|
model.to(DEVICE)
|
|
model.eval()
|
|
with torch.no_grad():
|
|
val_loader = torch.utils.data.DataLoader(val_ds, batch_size=64, collate_fn=dataloader.collate_to_cpu, num_workers=4)
|
|
results = []
|
|
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 // step, pred // step)
|
|
results.append((y.cpu(), y_pred.cpu()))
|
|
y_true = torch.cat([r[0] for r in results], dim=0)
|
|
y_pred = torch.cat([r[1] for r in results], dim=0)
|
|
np_y_all = np.concatenate(
|
|
[y_true.numpy().reshape(-1, 3), y_pred.numpy().reshape(-1, 3)], axis=1
|
|
)
|
|
np.savetxt(
|
|
f"predictions_{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.csv",
|
|
np_y_all,
|
|
delimiter=",",
|
|
header="true_x,true_y,true_z,pred_x,pred_y,pred_z",
|
|
comments="",
|
|
)
|