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)