def evaluate(args, model, loader, num_samples):
ade_outer, fde_outer = [], []
total_traj = 0
with torch.no_grad():
for batch_idx, data in enumerate(loader):
data = [tensor.cuda() for tensor in data]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, obs_goals,
pred_goals_gt, non_linear_ped, loss_mask, seq_start_end) = data
# goals one-hot encoding
obs_goals_ohe = to_goals_one_hot(obs_goals, args.g_dim).cuda()
# adj matrix for current batch
if args.adjacency_type == 0:
adj_out = compute_adjs(args, seq_start_end).cuda()
elif args.adjacency_type == 1:
adj_out = compute_adjs_distsim(
args, seq_start_end,
obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
elif args.adjacency_type == 2:
adj_out = compute_adjs_knnsim(
args, seq_start_end,
obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
kld, nll, ce, h = model(obs_traj, obs_traj_rel, obs_goals_ohe,
seq_start_end, adj_out)
ade, fde = [], []
total_traj += obs_traj.shape[1]
for _ in range(num_samples):
samples_rel = model.sample(args.pred_len, h, obs_traj[-1],
obs_goals_ohe[-1], seq_start_end)
samples = relative_to_abs(samples_rel, obs_traj[-1])
ade.append(
average_displacement_error(samples,
pred_traj_gt,
mode='raw'))
fde.append(
final_displacement_error(samples[-1, :, :],
pred_traj_gt[-1, :, :],
mode='raw'))
ade_sum = evaluate_helper(ade, seq_start_end)
fde_sum = evaluate_helper(fde, seq_start_end)
ade_outer.append(ade_sum)
fde_outer.append(fde_sum)
ade = sum(ade_outer) / (total_traj * args.pred_len)
fde = sum(fde_outer) / total_traj
return ade, fde
def validate(args, epch, valid_loader, model, warmup, lr_scheduler, writer):
test_loss = 0
kld_loss = 0
nll_loss = 0
cross_entropy_loss = 0
total_traj = 0
ade = 0
fde = 0
model.eval()
with torch.no_grad():
for batch_idx, data in enumerate(valid_loader):
data = [tensor.cuda() for tensor in data]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_rel_gt,
obs_goals, pred_goals_gt, seq_start_end) = data
# goals one-hot encoding
obs_goals_ohe = to_goals_one_hot(obs_goals, args.g_dim).cuda()
# adj matrix for current batch
if args.adjacency_type == 0:
adj_out = compute_adjs(args, seq_start_end).cuda()
elif args.adjacency_type == 1:
adj_out = compute_adjs_distsim(args, seq_start_end, obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
elif args.adjacency_type == 2:
adj_out = compute_adjs_knnsim(args, seq_start_end, obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
kld, nll, ce, h = model(obs_traj, obs_traj_rel, obs_goals_ohe, seq_start_end, adj_out)
loss = nll + (warmup[epch-1] * kld) + (ce * args.CE_weight)
test_loss += loss.item()
kld_loss += kld.item()
nll_loss += nll.item()
cross_entropy_loss += ce.item()
if batch_idx % args.print_every_batch == 0:
writer.add_scalar('validation/loss_items', loss.item(), epch) if not DEBUG else None
# predict trajectories from latest h; samples_rel shape=(pred_seq_len, n_agents, batch, xy)
samples_rel = model.sample(args.pred_len, h, obs_traj[-1], obs_goals_ohe[-1], seq_start_end)
samples = relative_to_abs(samples_rel, obs_traj[-1])
total_traj += samples.shape[1] # num_seqs
ade += average_displacement_error(samples, pred_traj_gt)
fde += final_displacement_error(samples[-1, :, :], pred_traj_gt[-1, :, :])
mean_loss = loss / len(valid_loader.dataset)
writer.add_scalar('validation/Loss', mean_loss, epch) if not DEBUG else None
mean_kld_loss = kld_loss / len(valid_loader.dataset)
mean_nll_loss = nll_loss / len(valid_loader.dataset)
mean_cross_entropy_loss = cross_entropy_loss / len(valid_loader.dataset)
writer.add_scalar('validation/KLD_loss', mean_kld_loss, epch) if not DEBUG else None
writer.add_scalar('validation/NLL_loss', mean_nll_loss, epch) if not DEBUG else None
writer.add_scalar('validation/CE_loss', mean_cross_entropy_loss, epch) if not DEBUG else None
# ADE
ade_val = ade / (total_traj * args.pred_len)
writer.add_scalar('validation/ADE', ade_val, epch) if not DEBUG else None
# FDE
fde_val = fde / total_traj
writer.add_scalar('validation/FDE', fde_val, epch) if not DEBUG else None
# eventually scheduler step
if args.lr_scheduler and epch > 100:
# let the metric settle for the first n epochs, then let the scheduler take care of the LR
lr_scheduler.step(ade_val)
return ade_val, fde_val
def train(args, epch, train_loader, model, warmup, optimizer, writer):
train_loss = 0
kld_loss = 0
nll_loss = 0
cross_entropy_loss = 0
start = time()
model.train()
for batch_idx, data in enumerate(train_loader):
data = [tensor.cuda() for tensor in data]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_rel_gt,
obs_goals, pred_goals_gt, seq_start_end) = data
#seq_start_end = seq_start_end.transpose(0, 1)
seq_len = len(obs_traj) + len(pred_traj_gt)
assert seq_len == args.obs_len + args.pred_len
# goals one-hot encoding
obs_goals_ohe = to_goals_one_hot(obs_goals, args.g_dim).cuda()
pred_goals_gt_ohe = to_goals_one_hot(pred_goals_gt, args.g_dim).cuda()
# adj matrix for current batch
if args.adjacency_type == 0:
adj_out = compute_adjs(args, seq_start_end).cuda()
elif args.adjacency_type == 1:
adj_out = compute_adjs_distsim(args, seq_start_end, obs_traj.detach().cpu(), pred_traj_gt.detach().cpu()).cuda()
elif args.adjacency_type == 2:
adj_out = compute_adjs_knnsim(args, seq_start_end, obs_traj.detach().cpu(), pred_traj_gt.detach().cpu()).cuda()
# during training we feed the entire trjs to the model
all_traj = torch.cat((obs_traj, pred_traj_gt), dim=0)
all_traj_rel = torch.cat((obs_traj_rel, pred_traj_rel_gt), dim=0)
all_goals_ohe = torch.cat((obs_goals_ohe, pred_goals_gt_ohe), dim=0)
optimizer.zero_grad()
kld, nll, ce, _ = model(all_traj, all_traj_rel, all_goals_ohe, seq_start_end, adj_out)
loss = (warmup[epch-1] * kld) + nll + (ce * args.CE_weight)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
optimizer.step()
train_loss += loss.item()
kld_loss += kld.item()
nll_loss += nll.item()
cross_entropy_loss += ce.item()
if batch_idx % args.print_every_batch == 0:
writer.add_scalar('train/loss_items', loss.item(), epch) if not DEBUG else None
avg_loss = train_loss / len(train_loader.dataset)
writer.add_scalar('train/Loss', avg_loss, epch) if not DEBUG else None
avg_kld_loss = kld_loss / len(train_loader.dataset)
avg_nll_loss = nll_loss / len(train_loader.dataset)
avg_cross_entropy_loss = cross_entropy_loss / len(train_loader.dataset)
writer.add_scalar('train/KLD_loss', avg_kld_loss, epch) if not DEBUG else None
writer.add_scalar('train/NLL_loss', avg_nll_loss, epch) if not DEBUG else None
writer.add_scalar('train/CE_loss', avg_cross_entropy_loss, epch) if not DEBUG else None
end = time()
elapsed = str(timedelta(seconds=(ceil(end - start))))
logging.info('Epoch [{}], time elapsed: {}'.format(epch, elapsed))
def test(args, epch, test_loader, model, warmup, writer):
test_loss = 0
kld_loss = 0
nll_loss = 0
cross_entropy_loss = 0
total_traj = 0
ade = 0
fde = 0
model.eval()
with torch.no_grad():
for batch_idx, data in enumerate(test_loader):
data = [tensor.cuda() for tensor in data]
(obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_rel_gt,
obs_goals, pred_goals_gt, seq_start_end) = data
# goals one-hot encoding
obs_goals_ohe = to_goals_one_hot(obs_goals, args.g_dim).cuda()
# adj matrix for current batch
if args.adjacency_type == 0:
adj_out = compute_adjs(args, seq_start_end).cuda()
elif args.adjacency_type == 1:
adj_out = compute_adjs_distsim(args, seq_start_end, obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
elif args.adjacency_type == 2:
adj_out = compute_adjs_knnsim(args, seq_start_end, obs_traj.detach().cpu(),
pred_traj_gt.detach().cpu()).cuda()
kld, nll, ce, h = model(obs_traj, obs_traj_rel, obs_goals_ohe, seq_start_end, adj_out)
loss = nll + (warmup[epch - 1] * kld) + (ce * args.CE_weight)
test_loss += loss.item()
kld_loss += kld.item()
nll_loss += nll.item()
cross_entropy_loss += ce.item()
if batch_idx % args.print_every_batch == 0:
writer.add_scalar('test/loss_items', loss.item(), epch) if not DEBUG else None
# predict trajectories from latest h; samples_rel shape=(pred_seq_len, n_agents, batch, xy)
samples_rel = model.sample(args.pred_len, h, obs_traj[-1], obs_goals_ohe[-1], seq_start_end)
samples = relative_to_abs(samples_rel, obs_traj[-1])
total_traj += samples.shape[1] # num_seqs
ade += average_displacement_error(samples, pred_traj_gt)
fde += final_displacement_error(samples[-1, :, :], pred_traj_gt[-1, :, :])
mean_loss = loss / len(test_loader.dataset)
writer.add_scalar('test/Loss', mean_loss, epch) if not DEBUG else None
mean_kld_loss = kld_loss / len(test_loader.dataset)
mean_nll_loss = nll_loss / len(test_loader.dataset)
mean_cross_entropy_loss = cross_entropy_loss / len(test_loader.dataset)
writer.add_scalar('test/KLD_loss', mean_kld_loss, epch) if not DEBUG else None
writer.add_scalar('test/NLL_loss', mean_nll_loss, epch) if not DEBUG else None
writer.add_scalar('test/CE_loss', mean_cross_entropy_loss, epch) if not DEBUG else None
# ADE
ade_val = ade / (total_traj * args.pred_len)
writer.add_scalar('test/ADE', ade_val, epch) if not DEBUG else None
# FDE
fde_val = fde / total_traj
writer.add_scalar('test/FDE', fde_val, epch) if not DEBUG else None
# plotting
if not DEBUG:
obs = obs_traj.cpu().numpy()
pred = samples.cpu().numpy()
pred_gt = pred_traj_gt.cpu().numpy()
plot_traj(obs, pred_gt, pred, seq_start_end, writer, epch)