-
Notifications
You must be signed in to change notification settings - Fork 81
/
test_pretrained_model.py
129 lines (100 loc) · 4.93 KB
/
test_pretrained_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import sys
from torch.utils.data import DataLoader
import argparse
import copy
sys.path.append("../utils/")
import matplotlib.pyplot as plt
import numpy as np
from models import *
from social_utils import *
import yaml
parser = argparse.ArgumentParser(description='PECNet')
parser.add_argument('--num_workers', '-nw', type=int, default=0)
parser.add_argument('--gpu_index', '-gi', type=int, default=0)
parser.add_argument('--load_file', '-lf', default="run7.pt")
parser.add_argument('--num_trajectories', '-nt', default=20) #number of trajectories to sample
parser.add_argument('--verbose', '-v', action='store_true')
parser.add_argument('--root_path', '-rp', default="./")
args = parser.parse_args()
dtype = torch.float64
torch.set_default_dtype(dtype)
device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu')
if torch.cuda.is_available():
torch.cuda.set_device(args.gpu_index)
print(device)
checkpoint = torch.load('../saved_models/{}'.format(args.load_file), map_location=device)
hyper_params = checkpoint["hyper_params"]
print(hyper_params)
def test(test_dataset, model, best_of_n = 1):
model.eval()
assert best_of_n >= 1 and type(best_of_n) == int
test_loss = 0
with torch.no_grad():
for i, (traj, mask, initial_pos) in enumerate(zip(test_dataset.trajectory_batches, test_dataset.mask_batches, test_dataset.initial_pos_batches)):
traj, mask, initial_pos = torch.DoubleTensor(traj).to(device), torch.DoubleTensor(mask).to(device), torch.DoubleTensor(initial_pos).to(device)
x = traj[:, :hyper_params["past_length"], :]
y = traj[:, hyper_params["past_length"]:, :]
y = y.cpu().numpy()
# reshape the data
x = x.view(-1, x.shape[1]*x.shape[2])
x = x.to(device)
future = y[:, :-1, :]
dest = y[:, -1, :]
all_l2_errors_dest = []
all_guesses = []
for index in range(best_of_n):
dest_recon = model.forward(x, initial_pos, device=device)
dest_recon = dest_recon.cpu().numpy()
all_guesses.append(dest_recon)
l2error_sample = np.linalg.norm(dest_recon - dest, axis = 1)
all_l2_errors_dest.append(l2error_sample)
all_l2_errors_dest = np.array(all_l2_errors_dest)
all_guesses = np.array(all_guesses)
# average error
l2error_avg_dest = np.mean(all_l2_errors_dest)
# choosing the best guess
indices = np.argmin(all_l2_errors_dest, axis = 0)
best_guess_dest = all_guesses[indices,np.arange(x.shape[0]), :]
# taking the minimum error out of all guess
l2error_dest = np.mean(np.min(all_l2_errors_dest, axis = 0))
# back to torch land
best_guess_dest = torch.DoubleTensor(best_guess_dest).to(device)
# using the best guess for interpolation
interpolated_future = model.predict(x, best_guess_dest, mask, initial_pos)
interpolated_future = interpolated_future.cpu().numpy()
best_guess_dest = best_guess_dest.cpu().numpy()
# final overall prediction
predicted_future = np.concatenate((interpolated_future, best_guess_dest), axis = 1)
predicted_future = np.reshape(predicted_future, (-1, hyper_params["future_length"], 2))
# ADE error
l2error_overall = np.mean(np.linalg.norm(y - predicted_future, axis = 2))
l2error_overall /= hyper_params["data_scale"]
l2error_dest /= hyper_params["data_scale"]
l2error_avg_dest /= hyper_params["data_scale"]
print('Test time error in destination best: {:0.3f} and mean: {:0.3f}'.format(l2error_dest, l2error_avg_dest))
print('Test time error overall (ADE) best: {:0.3f}'.format(l2error_overall))
return l2error_overall, l2error_dest, l2error_avg_dest
def main():
N = args.num_trajectories #number of generated trajectories
model = PECNet(hyper_params["enc_past_size"], hyper_params["enc_dest_size"], hyper_params["enc_latent_size"], hyper_params["dec_size"], hyper_params["predictor_hidden_size"], hyper_params['non_local_theta_size'], hyper_params['non_local_phi_size'], hyper_params['non_local_g_size'], hyper_params["fdim"], hyper_params["zdim"], hyper_params["nonlocal_pools"], hyper_params['non_local_dim'], hyper_params["sigma"], hyper_params["past_length"], hyper_params["future_length"], args.verbose)
model = model.double().to(device)
model.load_state_dict(checkpoint["model_state_dict"])
test_dataset = SocialDataset(set_name="test", b_size=hyper_params["test_b_size"], t_tresh=hyper_params["time_thresh"], d_tresh=hyper_params["dist_thresh"], verbose=args.verbose)
for traj in test_dataset.trajectory_batches:
traj -= traj[:, :1, :]
traj *= hyper_params["data_scale"]
#average ade/fde for k=20 (to account for variance in sampling)
num_samples = 150
average_ade, average_fde = 0, 0
for i in range(num_samples):
test_loss, final_point_loss_best, final_point_loss_avg = test(test_dataset, model, best_of_n = N)
average_ade += test_loss
average_fde += final_point_loss_best
print()
print("Average ADE:", average_ade/num_samples)
print("Average FDE:", average_fde/num_samples)
main()