def main(): parser = argparse.ArgumentParser( description='Train the individual Transformer model') parser.add_argument('--dataset_folder', type=str, default='datasets') parser.add_argument('--dataset_name', type=str, default='zara1') parser.add_argument('--obs', type=int, default=12) # size of history steps in frames parser.add_argument('--preds', type=int, default=8) # size of predicted trajectory in frames parser.add_argument('--point_dim', type=int, default=3) # number of dimensions (x,y,z) is 3 parser.add_argument('--emb_size', type=int, default=512) parser.add_argument('--heads', type=int, default=8) parser.add_argument('--layers', type=int, default=6) parser.add_argument('--dropout', type=float, default=0.1) parser.add_argument('--cpu', action='store_true') parser.add_argument('--val_size', type=int, default=0) parser.add_argument('--verbose', action='store_true') parser.add_argument('--max_epoch', type=int, default=1500) parser.add_argument('--batch_size', type=int, default=70) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train', action='store_true') parser.add_argument('--delim', type=str, default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--factor', type=float, default=1.) parser.add_argument('--save_step', type=int, default=1) parser.add_argument('--warmup', type=int, default=2) parser.add_argument('--evaluate', type=bool, default=True) parser.add_argument('--gen_pth', type=str) parser.add_argument('--crit_pth', type=str) parser.add_argument('--visual_step', type=int, default=10) parser.add_argument('--grad_penality', type=float, default=10) parser.add_argument('--crit_repeats', type=int, default=5) parser.add_argument('--lambda_recon', type=float, default=0.1) parser.add_argument('--z_dim', type=int, default=3) parser.add_argument('--stop_recon', type=int, default=2) args = parser.parse_args() model_name = args.name def mkdir(path): try: os.mkdir(path) except: pass paths = [ 'models', 'models/gen', 'models/crit', 'models/gan', f'models/gen/{args.name}', f'models/crit/{args.name}', f'models/gan/{args.name}', 'output', 'output/gan', f'output/gan/{args.name}' ] for path in paths: mkdir(path) log = SummaryWriter('logs/gan_%s' % model_name) device = torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device = torch.device("cpu") args.verbose = True ## creation of the dataloaders for train and validation if args.val_size == 0: train_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) else: train_dataset, val_dataset = baselineUtils.create_dataset( args.dataset_folder, args.dataset_name, args.val_size, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) test_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, eval=True, verbose=args.verbose) # import individual_TF # model=individual_TF.IndividualTF(3, 4, 4, N=args.layers, # d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout,mean=[0,0],std=[0,0]).to(device) tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) # optim = NoamOpt(args.emb_size, args.factor, len(tr_dl)*args.warmup, # torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) #optim=Adagrad(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01,lr_decay=0.001) epoch = 0 #mean=train_dataset[:]['src'][:,1:,2:4].mean((0,1)) mean = torch.cat((train_dataset[:]['src'][:, 1:, -3:], train_dataset[:]['trg'][:, :, -3:]), 1).mean((0, 1)) #std=train_dataset[:]['src'][:,1:,2:4].std((0,1)) std = torch.cat((train_dataset[:]['src'][:, 1:, -3:], train_dataset[:]['trg'][:, :, -3:]), 1).std((0, 1)) means = [] stds = [] for i in np.unique(train_dataset[:]['dataset']): ind = train_dataset[:]['dataset'] == i means.append( torch.cat((train_dataset[:]['src'][ind, 1:, -3:], train_dataset[:]['trg'][ind, :, -3:]), 1).mean((0, 1))) stds.append( torch.cat((train_dataset[:]['src'][ind, 1:, -3:], train_dataset[:]['trg'][ind, :, -3:]), 1).std((0, 1))) mean = torch.stack(means).mean(0) std = torch.stack(stds).mean(0) scipy.io.savemat(f'models/gan/{args.name}/norm.mat', { 'mean': mean.cpu().numpy(), 'std': std.cpu().numpy() }) from gan import Generator, Critic, get_gradient, gradient_penalty, get_crit_loss, get_gen_loss from tqdm import tqdm c_lambda = args.grad_penality crit_repeats = args.crit_repeats gen = Generator(args.obs - 1, args.preds, args.point_dim, args.point_dim, args.point_dim, z_dim=args.z_dim, N=args.layers, d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout, device=device).to(device) gen_opt = torch.optim.Adam(gen.parameters()) # gen_opt = NoamOpt(args.emb_size, args.factor, len(tr_dl)*args.warmup, # torch.optim.Adam(gen.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) crit = Critic(args.point_dim, args.obs - 1 + args.preds, N=args.layers, d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout, device=device).to(device) crit_opt = torch.optim.Adam(crit.parameters()) # crit_opt = NoamOpt(args.emb_size, args.factor, len(tr_dl)*args.warmup, # torch.optim.Adam(crit.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) if args.resume_train: gen.load_state_dict( torch.load(f'models/gen/{args.name}/{args.gen_pth}')) crit.load_state_dict( torch.load(f'models/crit/{args.name}/{args.crit_pth}')) cur_step = -1 for epoch in range(args.max_epoch): gen.train() crit.train() for id_b, batch in enumerate(tqdm(tr_dl, desc=f"Epoch {epoch}")): cur_step += 1 src = (batch['src'][:, 1:, -3:].to(device) - mean.to(device)) / std.to(device) tgt = (batch['trg'][:, :, -3:].to(device) - mean.to(device)) / std.to(device) batch_size = src.shape[0] mean_iteration_critic_loss = 0 for _ in range(crit_repeats): ### Update critic ### crit_opt.zero_grad() fake_noise = gen.sample_noise(batch_size) fake = gen(src, fake_noise) fake_seq = torch.cat((src, fake.detach()), dim=1) real_seq = torch.cat((src, tgt), dim=1) crit_fake_pred = crit(fake_seq) crit_real_pred = crit(real_seq) crit_loss = get_crit_loss( crit, src, tgt, fake.detach(), crit_fake_pred, crit_real_pred, c_lambda, args.lambda_recon if epoch < args.stop_recon else 0.) mean_iteration_critic_loss += crit_loss.item() / crit_repeats crit_loss.backward(retain_graph=True) crit_opt.step() log.add_scalar('Loss/train/crit', mean_iteration_critic_loss, cur_step) ### Update generator ### gen_opt.zero_grad() fake_noise_2 = gen.sample_noise(batch_size) fake_2 = gen(src, fake_noise_2) fake_2_seq = torch.cat((src, fake_2), dim=1) crit_fake_pred = crit(fake_2_seq) gen_loss = get_gen_loss( crit_fake_pred, fake_2, tgt, args.lambda_recon if epoch < args.stop_recon else 0.) gen_loss.backward() gen_opt.step() log.add_scalar('Loss/train/gen', gen_loss.item(), cur_step) if cur_step % args.visual_step == 0: scipy.io.savemat( f"output/gan/{args.name}/step_{cur_step:05}.mat", { 'input': batch['src'][:, 1:, :3].detach().cpu().numpy(), 'gt': batch['trg'][:, :, :3].detach().cpu().numpy(), 'pr': (fake_2 * std.to(device) + mean.to(device)).detach().cpu().numpy().cumsum(1) + batch['src'][:, -1:, :3].cpu().numpy() }) if epoch % args.save_step == 0: torch.save(gen.state_dict(), f'models/gen/{args.name}/{cur_step:05}.pth') torch.save(crit.state_dict(), f'models/crit/{args.name}/{cur_step:05}.pth')
def main(): parser = argparse.ArgumentParser( description='Train the individual Transformer model') parser.add_argument('--dataset_folder', type=str, default='datasets') parser.add_argument('--dataset_name', type=str, default='zara1') parser.add_argument('--obs', type=int, default=8) parser.add_argument('--preds', type=int, default=12) parser.add_argument('--emb_size', type=int, default=1024) parser.add_argument('--heads', type=int, default=8) parser.add_argument('--layers', type=int, default=6) parser.add_argument('--dropout', type=float, default=0.1) parser.add_argument('--cpu', action='store_true') parser.add_argument('--output_folder', type=str, default='Output') parser.add_argument('--val_size', type=int, default=50) parser.add_argument('--gpu_device', type=str, default="0") parser.add_argument('--verbose', action='store_true') parser.add_argument('--max_epoch', type=int, default=100) parser.add_argument('--batch_size', type=int, default=256) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train', action='store_true') parser.add_argument('--delim', type=str, default='\t') parser.add_argument('--name', type=str, default="zara1") args = parser.parse_args() model_name = args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/BERT') except: pass try: os.mkdir(f'models/BERT') except: pass try: os.mkdir(f'output/BERT/{args.name}') except: pass try: os.mkdir(f'models/BERT/{args.name}') except: pass log = SummaryWriter('logs/BERT_%s' % model_name) log.add_scalar('eval/mad', 0, 0) log.add_scalar('eval/fad', 0, 0) try: os.mkdir(args.name) except: pass device = torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device = torch.device("cpu") args.verbose = True ## creation of the dataloaders for train and validation train_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) test_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, eval=True, verbose=args.verbose) from transformers import BertTokenizer, BertModel, BertForMaskedLM, BertConfig, AdamW config = BertConfig(vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='relu', hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12) model = BertModel(config).to(device) from individual_TF import LinearEmbedding as NewEmbed, Generator as GeneratorTS a = NewEmbed(3, 768).to(device) model.set_input_embeddings(a) generator = GeneratorTS(768, 2).to(device) #model.set_output_embeddings(GeneratorTS(1024,2)) tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) optim = NoamOpt( 768, 0.1, len(tr_dl), torch.optim.Adam(list(a.parameters()) + list(model.parameters()) + list(generator.parameters()), lr=0, betas=(0.9, 0.98), eps=1e-9)) #optim=Adagrad(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01,lr_decay=0.001) epoch = 0 mean = train_dataset[:]['src'][:, :, 2:4].mean((0, 1)) * 0 std = train_dataset[:]['src'][:, :, 2:4].std((0, 1)) * 0 + 1 while epoch < args.max_epoch: epoch_loss = 0 model.train() for id_b, batch in enumerate(tr_dl): optim.optimizer.zero_grad() r = 0 rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) inp = ((batch['src'][:, :, 2:4] - mean) / std).to(device) inp = torch.matmul(inp, torch.from_numpy(rot_mat).float().to(device)) trg_masked = torch.zeros((inp.shape[0], args.preds, 2)).to(device) inp_cls = torch.ones(inp.shape[0], inp.shape[1], 1).to(device) trg_cls = torch.zeros(trg_masked.shape[0], trg_masked.shape[1], 1).to(device) inp_cat = torch.cat((inp, trg_masked), 1) cls_cat = torch.cat((inp_cls, trg_cls), 1) net_input = torch.cat((inp_cat, cls_cat), 2) position = torch.arange(0, net_input.shape[1]).repeat( inp.shape[0], 1).long().to(device) token = torch.zeros( (inp.shape[0], net_input.shape[1])).long().to(device) attention_mask = torch.ones( (inp.shape[0], net_input.shape[1])).long().to(device) out = model(input_ids=net_input, position_ids=position, token_type_ids=token, attention_mask=attention_mask) pred = generator(out[0]) loss = F.pairwise_distance( pred[:, :].contiguous().view(-1, 2), torch.matmul( torch.cat( (batch['src'][:, :, 2:4], batch['trg'][:, :, 2:4]), 1).contiguous().view(-1, 2).to(device), torch.from_numpy(rot_mat).float().to(device))).mean() loss.backward() optim.step() print("epoch %03i/%03i frame %04i / %04i loss: %7.4f" % (epoch, args.max_epoch, id_b, len(tr_dl), loss.item())) epoch_loss += loss.item() #sched.step() log.add_scalar('Loss/train', epoch_loss / len(tr_dl), epoch) with torch.no_grad(): model.eval() gt = [] pr = [] val_loss = 0 for batch in val_dl: inp = ((batch['src'][:, :, 2:4] - mean) / std).to(device) trg_masked = torch.zeros( (inp.shape[0], args.preds, 2)).to(device) inp_cls = torch.ones(inp.shape[0], inp.shape[1], 1).to(device) trg_cls = torch.zeros(trg_masked.shape[0], trg_masked.shape[1], 1).to(device) inp_cat = torch.cat((inp, trg_masked), 1) cls_cat = torch.cat((inp_cls, trg_cls), 1) net_input = torch.cat((inp_cat, cls_cat), 2) position = torch.arange(0, net_input.shape[1]).repeat( inp.shape[0], 1).long().to(device) token = torch.zeros( (inp.shape[0], net_input.shape[1])).long().to(device) attention_mask = torch.zeros( (inp.shape[0], net_input.shape[1])).long().to(device) out = model(input_ids=net_input, position_ids=position, token_type_ids=token, attention_mask=attention_mask) pred = generator(out[0]) loss = F.pairwise_distance( pred[:, :].contiguous().view(-1, 2), torch.cat( (batch['src'][:, :, 2:4], batch['trg'][:, :, 2:4]), 1).contiguous().view(-1, 2).to(device)).mean() val_loss += loss.item() gt_b = batch['trg'][:, :, 0:2] preds_tr_b = pred[:, args.obs:].cumsum(1).to( 'cpu').detach() + batch['src'][:, -1:, 0:2] gt.append(gt_b) pr.append(preds_tr_b) gt = np.concatenate(gt, 0) pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) log.add_scalar('validation/loss', val_loss / len(val_dl), epoch) log.add_scalar('validation/mad', mad, epoch) log.add_scalar('validation/fad', fad, epoch) model.eval() gt = [] pr = [] for batch in test_dl: inp = ((batch['src'][:, :, 2:4] - mean) / std).to(device) trg_masked = torch.zeros( (inp.shape[0], args.preds, 2)).to(device) inp_cls = torch.ones(inp.shape[0], inp.shape[1], 1).to(device) trg_cls = torch.zeros(trg_masked.shape[0], trg_masked.shape[1], 1).to(device) inp_cat = torch.cat((inp, trg_masked), 1) cls_cat = torch.cat((inp_cls, trg_cls), 1) net_input = torch.cat((inp_cat, cls_cat), 2) position = torch.arange(0, net_input.shape[1]).repeat( inp.shape[0], 1).long().to(device) token = torch.zeros( (inp.shape[0], net_input.shape[1])).long().to(device) attention_mask = torch.zeros( (inp.shape[0], net_input.shape[1])).long().to(device) out = model(input_ids=net_input, position_ids=position, token_type_ids=token, attention_mask=attention_mask) pred = generator(out[0]) gt_b = batch['trg'][:, :, 0:2] preds_tr_b = pred[:, args.obs:].cumsum(1).to( 'cpu').detach() + batch['src'][:, -1:, 0:2] gt.append(gt_b) pr.append(preds_tr_b) gt = np.concatenate(gt, 0) pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) torch.save(model.state_dict(), "models/BERT/%s/ep_%03i.pth" % (args.name, epoch)) torch.save(generator.state_dict(), "models/BERT/%s/gen_%03i.pth" % (args.name, epoch)) torch.save(a.state_dict(), "models/BERT/%s/emb_%03i.pth" % (args.name, epoch)) log.add_scalar('eval/mad', mad, epoch) log.add_scalar('eval/fad', fad, epoch) epoch += 1 ab = 1
def main(dataset_name, model_layers, emb_size, heads, obs=8, preds=12, dropout=0.): device = torch.device("cuda") model = individual_TF.IndividualTF(2, 3, 3, N=model_layers, d_model=emb_size, d_ff=2048, heads=heads, dropout=dropout, mean=[0, 0], std=[0, 0]).to(device) model.load_state_dict( torch.load(f'models/Individual/{dataset_name}/model.pth')) mean_std = scipy.io.loadmat(f'models/Individual/{dataset_name}/norm.mat') model.eval() inference_set, _ = baselineUtils.create_dataset("datasets", args.dataset_name, 0, obs, preds, verbose=True, inference=True) inference_dl = torch.utils.data.DataLoader(inference_set, batch_size=1, shuffle=False, num_workers=0) model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for id_b, batch in enumerate(inference_dl): inp_.append(batch['src']) gt.append(batch['trg'][:, :, 0:2]) frames.append(batch['frames']) peds.append(batch['peds']) dt.append(batch['dataset']) inp = (batch['src'][:, 1:, 2:4].to(device) - torch.tensor(mean_std["mean"], device=device)) / torch.tensor( mean_std["std"], device=device) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat( inp.shape[0], 1, 1).to(device) dec_inp = start_of_seq for i in range(12): trg_att = subsequent_mask(dec_inp.shape[1]).repeat( dec_inp.shape[0], 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp = torch.cat((dec_inp, out[:, -1:, :]), 1) preds_tr_b = (dec_inp[:, 1:, 0:2] * torch.tensor(mean_std["std"], device=device) + torch.tensor(mean_std["mean"], device=device)).cpu().detach().numpy().cumsum(1) + \ batch['src'][:, -1:, 0:2].cpu().numpy() pr.append(preds_tr_b) print("inference: batch %04i / %04i" % (id_b, len(inference_dl))) peds = np.concatenate(peds, 0) frames = np.concatenate(frames, 0) dt = np.concatenate(dt, 0) gt = np.concatenate(gt, 0) dt_names = inference_set.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) print(gt[0], pr[0]) print(mad, fad)
def main(): parser = argparse.ArgumentParser( description='Train the individual Transformer model') parser.add_argument('--dataset_folder', type=str, default='datasets') parser.add_argument('--dataset_name', type=str, default='zara1') parser.add_argument('--obs', type=int, default=8) parser.add_argument('--preds', type=int, default=12) parser.add_argument('--emb_size', type=int, default=512) parser.add_argument('--heads', type=int, default=8) parser.add_argument('--layers', type=int, default=6) parser.add_argument('--dropout', type=float, default=0.1) parser.add_argument('--cpu', action='store_true') parser.add_argument('--val_size', type=int, default=0) parser.add_argument('--verbose', action='store_true') parser.add_argument('--max_epoch', type=int, default=1500) parser.add_argument('--batch_size', type=int, default=70) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train', action='store_true') parser.add_argument('--delim', type=str, default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--factor', type=float, default=1.) parser.add_argument('--save_step', type=int, default=1) parser.add_argument('--warmup', type=int, default=10) parser.add_argument('--evaluate', type=bool, default=True) args = parser.parse_args() model_name = args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/Individual') except: pass try: os.mkdir(f'models/Individual') except: pass try: os.mkdir(f'output/Individual/{args.name}') except: pass try: os.mkdir(f'models/Individual/{args.name}') except: pass log = SummaryWriter('logs/Ind_%s' % model_name) log.add_scalar('eval/mad', 0, 0) log.add_scalar('eval/fad', 0, 0) device = torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device = torch.device("cpu") args.verbose = True ## creation of the dataloaders for train and validation if args.val_size == 0: train_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) else: train_dataset, val_dataset = baselineUtils.create_dataset( args.dataset_folder, args.dataset_name, args.val_size, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) test_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, eval=True, verbose=args.verbose) import individual_TF model = individual_TF.IndividualTF(2, 3, 3, N=args.layers, d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout, mean=[0, 0], std=[0, 0]).to(device) tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) optim = NoamOpt( args.emb_size, args.factor, len(tr_dl) * args.warmup, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) #optim=Adagrad(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01,lr_decay=0.001) epoch = 0 #mean=train_dataset[:]['src'][:,1:,2:4].mean((0,1)) mean = torch.cat((train_dataset[:]['src'][:, 1:, 2:4], train_dataset[:]['trg'][:, :, 2:4]), 1).mean((0, 1)) #std=train_dataset[:]['src'][:,1:,2:4].std((0,1)) std = torch.cat((train_dataset[:]['src'][:, 1:, 2:4], train_dataset[:]['trg'][:, :, 2:4]), 1).std((0, 1)) means = [] stds = [] for i in np.unique(train_dataset[:]['dataset']): ind = train_dataset[:]['dataset'] == i means.append( torch.cat((train_dataset[:]['src'][ind, 1:, 2:4], train_dataset[:]['trg'][ind, :, 2:4]), 1).mean((0, 1))) stds.append( torch.cat((train_dataset[:]['src'][ind, 1:, 2:4], train_dataset[:]['trg'][ind, :, 2:4]), 1).std((0, 1))) mean = torch.stack(means).mean(0) std = torch.stack(stds).mean(0) scipy.io.savemat(f'models/Individual/{args.name}/norm.mat', { 'mean': mean.cpu().numpy(), 'std': std.cpu().numpy() }) while epoch < args.max_epoch: epoch_loss = 0 model.train() for id_b, batch in enumerate(tr_dl): optim.optimizer.zero_grad() #将所有variable的grad设置为0 inp = (batch['src'][:, 1:, 2:4].to(device) - mean.to(device)) / std.to(device) target = (batch['trg'][:, :-1, 2:4].to(device) - mean.to(device)) / std.to(device) target_c = torch.zeros( (target.shape[0], target.shape[1], 1)).to(device) target = torch.cat((target, target_c), -1) start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat( target.shape[0], 1, 1).to(device) dec_inp = torch.cat((start_of_seq, target), 1) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(dec_inp.shape[1]).repeat( dec_inp.shape[0], 1, 1).to(device) pred = model(inp, dec_inp, src_att, trg_att) loss = F.pairwise_distance( pred[:, :, 0:2].contiguous().view(-1, 2), ((batch['trg'][:, :, 2:4].to(device) - mean.to(device)) / std.to(device)).contiguous().view( -1, 2).to(device)).mean() + torch.mean( torch.abs(pred[:, :, 2])) loss.backward() optim.step() # 更新variable的grad print("train epoch %03i/%03i batch %04i / %04i loss: %7.4f" % (epoch, args.max_epoch, id_b, len(tr_dl), loss.item())) epoch_loss += loss.item() #sched.step() log.add_scalar('Loss/train', epoch_loss / len(tr_dl), epoch) with torch.no_grad(): model.eval() # 不更新梯度和batchnormalize val_loss = 0 step = 0 model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for id_b, batch in enumerate(val_dl): inp_.append(batch['src']) gt.append(batch['trg'][:, :, 0:2]) frames.append(batch['frames']) peds.append(batch['peds']) dt.append(batch['dataset']) inp = (batch['src'][:, 1:, 2:4].to(device) - mean.to(device)) / std.to(device) src_att = torch.ones( (inp.shape[0], 1, inp.shape[1])).to(device) start_of_seq = torch.Tensor( [0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat(inp.shape[0], 1, 1).to(device) dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat( dec_inp.shape[0], 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp = torch.cat((dec_inp, out[:, -1:, :]), 1) preds_tr_b = (dec_inp[:, 1:, 0:2] * std.to(device) + mean.to(device)).cpu().numpy().cumsum( 1) + batch['src'][:, -1:, 0:2].cpu().numpy() pr.append(preds_tr_b) print("val epoch %03i/%03i batch %04i / %04i" % (epoch, args.max_epoch, id_b, len(val_dl))) peds = np.concatenate(peds, 0) frames = np.concatenate(frames, 0) dt = np.concatenate(dt, 0) gt = np.concatenate(gt, 0) dt_names = test_dataset.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) log.add_scalar('validation/MAD', mad, epoch) log.add_scalar('validation/FAD', fad, epoch) if args.evaluate: model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for id_b, batch in enumerate(test_dl): inp_.append(batch['src']) gt.append(batch['trg'][:, :, 0:2]) frames.append(batch['frames']) peds.append(batch['peds']) dt.append(batch['dataset']) inp = (batch['src'][:, 1:, 2:4].to(device) - mean.to(device)) / std.to(device) src_att = torch.ones( (inp.shape[0], 1, inp.shape[1])).to(device) start_of_seq = torch.Tensor([ 0, 0, 1 ]).unsqueeze(0).unsqueeze(1).repeat(inp.shape[0], 1, 1).to(device) dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat( dec_inp.shape[0], 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp = torch.cat((dec_inp, out[:, -1:, :]), 1) preds_tr_b = (dec_inp[:, 1:, 0:2] * std.to(device) + mean.to(device)).cpu().numpy().cumsum( 1) + batch['src'][:, -1:, 0:2].cpu().numpy() pr.append(preds_tr_b) print("test epoch %03i/%03i batch %04i / %04i" % (epoch, args.max_epoch, id_b, len(test_dl))) peds = np.concatenate(peds, 0) frames = np.concatenate(frames, 0) dt = np.concatenate(dt, 0) gt = np.concatenate(gt, 0) dt_names = test_dataset.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) log.add_scalar('eval/DET_mad', mad, epoch) log.add_scalar('eval/DET_fad', fad, epoch) # log.add_scalar('eval/DET_mad', mad, epoch) # log.add_scalar('eval/DET_fad', fad, epoch) scipy.io.savemat( f"output/Individual/{args.name}/det_{epoch}.mat", { 'input': inp, 'gt': gt, 'pr': pr, 'peds': peds, 'frames': frames, 'dt': dt, 'dt_names': dt_names }) if epoch % args.save_step == 0: torch.save(model.state_dict(), f'models/Individual/{args.name}/{epoch:05d}.pth') epoch += 1 ab = 1
def main(): parser=argparse.ArgumentParser(description='Train the individual Transformer model') parser.add_argument('--dataset_folder',type=str,default='datasets') parser.add_argument('--dataset_name',type=str,default='zara1') parser.add_argument('--obs',type=int,default=8) parser.add_argument('--preds',type=int,default=12) parser.add_argument('--emb_size',type=int,default=512) parser.add_argument('--heads',type=int, default=8) parser.add_argument('--layers',type=int,default=6) parser.add_argument('--cpu',action='store_true') parser.add_argument('--verbose',action='store_true') parser.add_argument('--batch_size',type=int,default=256) parser.add_argument('--delim',type=str,default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--epoch',type=str,default="00001") parser.add_argument('--num_samples', type=int, default="20") args=parser.parse_args() model_name=args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/QuantizedTF') except: pass try: os.mkdir(f'models/QuantizedTF') except: pass try: os.mkdir(f'output/QuantizedTF/{args.name}') except: pass try: os.mkdir(f'models/QuantizedTF/{args.name}') except: pass #log=SummaryWriter('logs/%s'%model_name) # log.add_scalar('eval/mad', 0, 0) # log.add_scalar('eval/fad', 0, 0) device=torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device=torch.device("cpu") args.verbose=True ## creation of the dataloaders for train and validation test_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=False,eval=True,verbose=args.verbose) mat = scipy.io.loadmat(os.path.join(args.dataset_folder, args.dataset_name, "clusters.mat")) clusters=mat['centroids'] model=quantized_TF.QuantizedTF(clusters.shape[0], clusters.shape[0]+1, clusters.shape[0], N=args.layers, d_model=args.emb_size, d_ff=1024, h=args.heads).to(device) model.load_state_dict(torch.load(f'models/QuantizedTF/{args.name}/{args.epoch}.pth')) model.to(device) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) # DETERMINISTIC MODE with torch.no_grad(): model.eval() gt=[] pr=[] inp_=[] peds=[] frames=[] dt=[] for id_b,batch in enumerate(test_dl): print(f"batch {id_b:03d}/{len(test_dl)}") peds.append(batch['peds']) frames.append(batch['frames']) dt.append(batch['dataset']) scale = np.random.uniform(0.5, 2) # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] gt_b = batch['trg'][:, :, 0:2] inp = torch.tensor( scipy.spatial.distance.cdist(speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to( device) src_att = torch.ones((inp.shape[0], 1,inp.shape[1])).to(device) start_of_seq = torch.tensor([clusters.shape[0]]).repeat(n_in_batch).unsqueeze(1).to(device) dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat(n_in_batch, 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp=torch.cat((dec_inp,out[:,-1:].argmax(dim=2)),1) preds_tr_b=clusters[dec_inp[:,1:].cpu().numpy()].cumsum(1)+batch['src'][:,-1:,0:2].cpu().numpy() gt.append(gt_b) pr.append(preds_tr_b) peds=np.concatenate(peds,0) frames=np.concatenate(frames,0) dt=np.concatenate(dt,0) gt=np.concatenate(gt,0) dt_names=test_dataset.data['dataset_name'] pr=np.concatenate(pr,0) mad,fad,errs=baselineUtils.distance_metrics(gt,pr) #log.add_scalar('eval/DET_mad', mad, epoch) #log.add_scalar('eval/DET_fad', fad, epoch) scipy.io.savemat(f"output/QuantizedTF/{args.name}/MM_deterministic.mat",{'input':inp,'gt':gt,'pr':pr,'peds':peds,'frames':frames,'dt':dt,'dt_names':dt_names}) print("Determinitic:") print("mad: %6.3f"%mad) print("fad: %6.3f" % fad) # MULTI MODALITY num_samples=args.num_samples model.eval() gt=[] pr_all={} inp_=[] peds=[] frames=[] dt=[] for sam in range(num_samples): pr_all[sam]=[] for id_b,batch in enumerate(test_dl): print(f"batch {id_b:03d}/{len(test_dl)}") peds.append(batch['peds']) frames.append(batch['frames']) dt.append(batch['dataset']) scale = np.random.uniform(0.5, 2) # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] gt_b = batch['trg'][:, :, 0:2] gt.append(gt_b) inp__=batch['src'][:,:,0:2] inp_.append(inp__) inp = torch.tensor( scipy.spatial.distance.cdist(speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to( device) src_att = torch.ones((inp.shape[0], 1,inp.shape[1])).to(device) start_of_seq = torch.tensor([clusters.shape[0]]).repeat(n_in_batch).unsqueeze(1).to(device) for sam in range(num_samples): dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat(n_in_batch, 1, 1).to(device) out = model.predict(inp, dec_inp, src_att, trg_att) h=out[:,-1] dec_inp=torch.cat((dec_inp,torch.multinomial(h,1)),1) preds_tr_b=clusters[dec_inp[:,1:].cpu().numpy()].cumsum(1)+batch['src'][:,-1:,0:2].cpu().numpy() pr_all[sam].append(preds_tr_b) peds=np.concatenate(peds,0) frames=np.concatenate(frames,0) dt=np.concatenate(dt,0) gt=np.concatenate(gt,0) dt_names=test_dataset.data['dataset_name'] #pr=np.concatenate(pr,0) inp=np.concatenate(inp_,0) samp = {} for k in pr_all.keys(): samp[k] = {} samp[k]['pr'] = np.concatenate(pr_all[k], 0) samp[k]['mad'], samp[k]['fad'], samp[k]['err'] = baselineUtils.distance_metrics(gt, samp[k]['pr']) ev = [samp[i]['err'] for i in range(num_samples)] e20 = np.stack(ev, -1) mad_samp=e20.mean(1).min(-1).mean() fad_samp=e20[:,-1].min(-1).mean() #mad,fad,errs=baselineUtils.distance_metrics(gt,pr) #log.add_scalar('eval/MM_mad', mad_samp, epoch) #log.add_scalar('eval/MM_fad', fad_samp, epoch) preds_all_fin=np.stack(list([samp[i]['pr'] for i in range(num_samples)]),-1) scipy.io.savemat(f"output/QuantizedTF/{args.name}/MM_{num_samples}.mat",{'input':inp,'gt':gt,'pr':preds_all_fin,'peds':peds,'frames':frames,'dt':dt,'dt_names':dt_names}) print("Determinitic:") print("mad: %6.3f"%mad) print("fad: %6.3f" % fad) print("Multimodality:") print("mad: %6.3f"%mad_samp) print("fad: %6.3f" % fad_samp)
def main(): parser=argparse.ArgumentParser(description='Train the individual Transformer model') parser.add_argument('--dataset_folder',type=str,default='datasets') parser.add_argument('--dataset_name',type=str,default='zara1') parser.add_argument('--obs',type=int,default=8) parser.add_argument('--preds',type=int,default=12) parser.add_argument('--emb_size',type=int,default=512) parser.add_argument('--heads',type=int, default=8) parser.add_argument('--layers',type=int,default=6) parser.add_argument('--dropout',type=float,default=0.1) parser.add_argument('--cpu',action='store_true') parser.add_argument('--val_size',type=int, default=0) parser.add_argument('--verbose',action='store_true') parser.add_argument('--max_epoch',type=int, default=1500) parser.add_argument('--batch_size',type=int,default=70) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train',action='store_true') parser.add_argument('--delim',type=str,default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--factor', type=float, default=1.) parser.add_argument('--save_step', type=int, default=1) parser.add_argument('--warmup', type=int, default=10) parser.add_argument('--evaluate', type=bool, default=True) parser.add_argument('--model_pth', type=str) args=parser.parse_args() model_name=args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/Individual') except: pass try: os.mkdir(f'models/Individual') except: pass try: os.mkdir(f'output/Individual/{args.name}') except: pass try: os.mkdir(f'models/Individual/{args.name}') except: pass #log=SummaryWriter('logs/Ind_%s'%model_name) #log.add_scalar('eval/mad', 0, 0) #log.add_scalar('eval/fad', 0, 0) device=torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device=torch.device("cpu") args.verbose=True if args.val_size==0: train_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=True,verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) else: train_dataset, val_dataset = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, args.val_size,args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) test_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=False,eval=True,verbose=args.verbose) import individual_TF model=individual_TF.IndividualTF(2, 3, 3, N=args.layers, d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout,mean=[0,0],std=[0,0]).to(device) model.load_state_dict(torch.load(f'models/Individual/my_data_train/00600.pth')) #tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) #val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #mean=train_dataset[:]['src'][:,1:,2:4].mean((0,1)) #mean=torch.cat((train_dataset[:]['src'][:,1:,2:4],train_dataset[:]['trg'][:,:,2:4]),1).mean((0,1)) #std=train_dataset[:]['src'][:,1:,2:4].std((0,1)) #std=torch.cat((train_dataset[:]['src'][:,1:,2:4],train_dataset[:]['trg'][:,:,2:4]),1).std((0,1)) #means=[] #stds=[] #for i in np.unique(train_dataset[:]['dataset']): # ind=train_dataset[:]['dataset']==i # means.append(torch.cat((train_dataset[:]['src'][ind, 1:, 2:4], train_dataset[:]['trg'][ind, :, 2:4]), 1).mean((0, 1))) # stds.append( # torch.cat((train_dataset[:]['src'][ind, 1:, 2:4], train_dataset[:]['trg'][ind, :, 2:4]), 1).std((0, 1))) #mean=torch.stack(means).mean(0) #std=torch.stack(stds).mean(0) model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for id_b,batch in enumerate(test_dl): inp_.append(batch['src']) gt.append(batch['trg'][:,:,0:2]) frames.append(batch['frames']) peds.append(batch['peds']) dt.append(batch['dataset']) inp = batch['src'][:, 1:, 2:4].to(device) #- mean.to(device)) / std.to(device) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat(inp.shape[0], 1, 1).to( device) dec_inp=start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat(dec_inp.shape[0], 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp=torch.cat((dec_inp,out[:,-1:,:]),1) preds_tr_b=(dec_inp[:,1:,0:2]).cpu().detach().numpy().cumsum(1)+batch['src'][:,-1:,0:2].cpu().detach().numpy() pr.append(preds_tr_b) # print("test epoch %03i/%03i batch %04i / %04i" % ( # epoch, args.max_epoch, id_b, len(test_dl))) peds = np.concatenate(peds, 0) frames = np.concatenate(frames, 0) dt = np.concatenate(dt, 0) gt = np.concatenate(gt, 0) dt_names = test_dataset.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) #print(frames) #print(dt.shape) #print(dt) #print(gt[1]) #print(pr[1]) print("mad %f fad %f"%(mad,fad)) for i in range(pr.shape[0]): pathin = 'c_1 frames/c_1_' pathout = '5_1 frames_out/' img = cv2.imread(pathin+str(frames[i][8])+'.jpg') cg = (0,255,0) # green cp = (0,0,255) # red #print(gt[i]) #print(pr[i]) for j in range(12): gp = (int(gt[i,j,0]*1920),int(gt[i,j,1]*1080)) pp = (int(pr[i,j,0]*1920),int(pr[i,j,1]*1080)) img = cv2.circle(img,gp,3,cg,-1) img = cv2.circle(img,pp,3,cp,-1) #print(gp) #print(pp) #print(frames[i][8]) cv2.imwrite(pathout+str(frames[i][8])+'.jpg',img)
def main(): parser=argparse.ArgumentParser(description='Train the individual Transformer model') parser.add_argument('--dataset_folder',type=str,default='datasets') parser.add_argument('--dataset_name',type=str,default='eth') parser.add_argument('--obs',type=int,default=8) parser.add_argument('--preds',type=int,default=12) parser.add_argument('--emb_size',type=int,default=1024) parser.add_argument('--heads',type=int, default=8) parser.add_argument('--layers',type=int,default=6) parser.add_argument('--dropout',type=float,default=0.1) parser.add_argument('--cpu',action='store_true') parser.add_argument('--output_folder',type=str,default='Output') parser.add_argument('--val_size',type=int, default=50) parser.add_argument('--verbose',action='store_true') parser.add_argument('--max_epoch',type=int, default=100) parser.add_argument('--batch_size',type=int,default=256) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train',action='store_true') parser.add_argument('--delim',type=str,default='\t') parser.add_argument('--name', type=str, default="test_rot") parser.add_argument('--num_clusters', type=int, default=1000) parser.add_argument('--max_samples', type=int, default=200000) parser.add_argument('--scale', type=bool, default="True") parser.add_argument('--rot', type=bool, default="False") parser.add_argument('--axes_lim', type=int, default=2) args=parser.parse_args() model_name=args.name device=torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device=torch.device("cpu") args.verbose=True outdir=f'{args.dataset_name}_{args.num_clusters}_{args.max_samples}_scale{args.scale}_rot{args.rot}' try: os.mkdir(outdir) except: pass ## creation of the dataloaders for train and validation train_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=True,verbose=args.verbose) # val_dataset, _ = baselineUtils.create_train_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, # args.preds, delim=args.delim, train=False, # verbose=args.verbose) test_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=False,eval=True,verbose=args.verbose) tr = train_dataset[:]['src'][:, 1:, 2:4].reshape(-1, 2) pr = train_dataset[:]['trg'][:, :, 2:4].reshape(-1, 2) t = torch.cat((tr, pr), 0) t= t.cpu().numpy() if args.scale: s=[1,0.7,0.5,1.5,2] t2=[] for i in s: t2.append(t*i) t=np.concatenate(t2,0) plt.figure() plt.scatter(t[:,0],t[:,1]) plt.xlim([-args.axes_lim,args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/train_distribution.png') plt.close() te = test_dataset[:]['src'][:, 1:, 2:4].reshape(-1, 2) test=te if args.dataset_name!="trajnet": pe = test_dataset[:]['trg'][:, :, 2:4].reshape(-1, 2) test= torch.cat((te, pe), 0) test = test.cpu().numpy() plt.figure() plt.scatter(t[:,0],t[:,1],) plt.scatter(test[:,0],test[:,1]) plt.xlim([-args.axes_lim,args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/test_distribution.png') plt.close() plt.figure() plt.scatter(test[:, 0], test[:, 1],c=['orange']*test.shape[0]) plt.scatter(t[:, 0], t[:, 1],c=['blue']*t.shape[0], alpha=0.01) plt.xlim([-args.axes_lim, args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/test_distribution_alpha.png') plt.close() args.max_samples=min(args.max_samples,t.shape[0]) ind = np.random.choice(t.shape[0], args.max_samples, replace=False) ended,cluster_index, center = kmeans.lloyd(t[ind], args.num_clusters,tol=1e-3) if ended==False: print("kmeans è crashato") return 0 plt.figure() plt.set_cmap('prism') plt.scatter(t[ind, 0], t[ind, 1], c=cluster_index) plt.xlim([-args.axes_lim, args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/cluster_train_distribution_limited_data.png') plt.close() import scipy.spatial.distance as dist dist_tr = dist.cdist(t, center) dist_tr.min(axis=1).mean() dist_test = dist.cdist(test, center) dist_test.min(axis=1).mean() plt.figure() plt.set_cmap('prism') plt.scatter(t[:, 0], t[:, 1], c=dist_tr.argmin(1)) plt.xlim([-args.axes_lim, args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/cluster_train_distribution_full_data.png') plt.close() plt.figure() plt.set_cmap('prism') plt.scatter(test[:, 0], test[:, 1], c=dist_test.argmin(1)) plt.xlim([-args.axes_lim, args.axes_lim]) plt.ylim([-args.axes_lim, args.axes_lim]) plt.savefig(f'{outdir}/cluster_test_distribution_full_data.png') plt.close() plt.boxplot([dist_tr.min(axis=1),dist_test.min(axis=1)]) plt.savefig(f'{outdir}/residuals.png') plt.close() cluster_count_tr= np.bincount(dist_tr.argmin(1)) cluster_count_te= np.bincount(dist_test.argmin(1)) plt.figure() plt.bar(np.arange(len(cluster_count_tr)),cluster_count_tr) plt.savefig(f'{outdir}/cluster_use_tr.png') plt.close() plt.figure() plt.bar(np.arange(len(cluster_count_te)),cluster_count_te) plt.savefig(f'{outdir}/cluster_use_te.png') plt.close() with open(f'{outdir}/stats.json','w') as f: j={"train":{ "mean":dist_tr.min(axis=1).mean(), "std": dist_tr.min(axis=1).std(), "n_points":int(dist_tr.shape[0]), "avg_clust_presence":cluster_count_tr.mean(), "min_clust_presence": int(cluster_count_tr.min()), "max_clust_presence": int(cluster_count_tr.max()), "counts": cluster_count_tr.tolist() }, "test":{ "mean": dist_test.min(axis=1).mean(), "std": dist_test.min(axis=1).std(), "n_points":int(dist_test.shape[0]), "avg_clust_presence": cluster_count_te.mean(), "min_clust_presence": int(cluster_count_te.min()), "max_clust_presence": int(cluster_count_te.max()), "counts": cluster_count_te.tolist() } } json.dump(j,f) scipy.io.savemat(f'{outdir}/clusters.mat',{'centroids':center,"counts":cluster_count_tr})
def main(): parser=argparse.ArgumentParser(description='Train the individual Transformer model') parser.add_argument('--dataset_folder',type=str,default='datasets') parser.add_argument('--dataset_name',type=str,default='zara1') parser.add_argument('--obs',type=int,default=8) parser.add_argument('--preds',type=int,default=12) parser.add_argument('--emb_size',type=int,default=512) parser.add_argument('--heads',type=int, default=8) parser.add_argument('--layers',type=int,default=6) parser.add_argument('--dropout',type=float,default=0.1) parser.add_argument('--cpu',action='store_true') parser.add_argument('--val_size',type=int, default=0) parser.add_argument('--verbose',action='store_true') parser.add_argument('--max_epoch',type=int, default=1500) parser.add_argument('--batch_size',type=int,default=1) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train',action='store_true') parser.add_argument('--delim',type=str,default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--factor', type=float, default=1.) parser.add_argument('--save_step', type=int, default=1) parser.add_argument('--warmup', type=int, default=10) parser.add_argument('--evaluate', type=bool, default=True) parser.add_argument('--model_pth', type=str) args=parser.parse_args() model_name=args.name #device=torch.device("cuda") #if args.cpu or not torch.cuda.is_available(): device=torch.device("cpu") args.verbose=True test_dataset,_ = baselineUtils.create_dataset(args.dataset_folder,args.dataset_name,0,args.obs,args.preds,delim=args.delim,train=False,eval=True,verbose=args.verbose) import individual_TF model=individual_TF.IndividualTF(2, 3, 3, N=args.layers, d_model=args.emb_size, d_ff=2048, h=args.heads, dropout=args.dropout,mean=[0,0],std=[0,0]).to(device) model.load_state_dict(torch.load(f'models/Individual/my_data_train/00013.pth')) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for id_b,batch in enumerate(test_dl): #print(batch['src'].shape) #inp_.append(batch['src']) gt.append(batch['trg'][:,:,0:2]) #frames.append(batch['frames']) #peds.append(batch['peds']) #dt.append(batch['dataset']) inp = batch['src'][:, 1:, 2:4].to(device) #- mean.to(device)) / std.to(device) #print(inp.shape) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) start_of_seq = torch.Tensor([0, 0, 1]).unsqueeze(0).unsqueeze(1).repeat(inp.shape[0], 1, 1).to( device) # print("start of seq") # print(start_of_seq[0]) dec_inp=start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat(dec_inp.shape[0], 1, 1).to(device) # print("src_att shape") # print(src_att.shape) # print("trg_att shape") # print(trg_att.shape) out = model(inp, dec_inp, src_att, trg_att) # print("out shape") # print(out.shape) # print("-----------") dec_inp=torch.cat((dec_inp,out[:,-1:,:]),1) print("batch['src']") print(batch['src'].shape) preds_tr_b=(dec_inp[:,1:,0:2]).cpu().detach().numpy().cumsum(1)+batch['src'][:,-1:,0:2].cpu().detach().numpy() #print(preds_tr_b[1]) pr.append(preds_tr_b) # print("test epoch %03i/%03i batch %04i / %04i" % ( # epoch, args.max_epoch, id_b, len(test_dl))) gt = np.concatenate(gt, 0) #dt_names = test_dataset.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr)
def main(): parser = argparse.ArgumentParser( description='Train the individual Transformer model') parser.add_argument('--dataset_folder', type=str, default='datasets') parser.add_argument('--dataset_name', type=str, default='zara1') parser.add_argument('--obs', type=int, default=8) parser.add_argument('--preds', type=int, default=12) parser.add_argument('--emb_size', type=int, default=512) parser.add_argument('--heads', type=int, default=8) parser.add_argument('--layers', type=int, default=6) parser.add_argument('--dropout', type=float, default=0.1) parser.add_argument('--cpu', action='store_true') parser.add_argument('--output_folder', type=str, default='Output') parser.add_argument('--val_size', type=int, default=0) parser.add_argument('--gpu_device', type=str, default="0") parser.add_argument('--verbose', action='store_true') parser.add_argument('--max_epoch', type=int, default=100) parser.add_argument('--batch_size', type=int, default=100) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train', action='store_true') parser.add_argument('--delim', type=str, default='\t') parser.add_argument('--name', type=str, default="zara1") parser.add_argument('--factor', type=float, default=1.) parser.add_argument('--evaluate', type=bool, default=True) parser.add_argument('--save_step', type=int, default=1) args = parser.parse_args() model_name = args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/QuantizedTF') except: pass try: os.mkdir(f'models/QuantizedTF') except: pass try: os.mkdir(f'output/QuantizedTF/{args.name}') except: pass try: os.mkdir(f'models/QuantizedTF/{args.name}') except: pass log = SummaryWriter('logs/%s' % model_name) #os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_device device = torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device = torch.device("cpu") args.verbose = True ## creation of the dataloaders for train and validation if args.val_size == 0: train_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) else: train_dataset, val_dataset = baselineUtils.create_dataset( args.dataset_folder, args.dataset_name, args.val_size, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) test_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, eval=True, verbose=args.verbose) mat = scipy.io.loadmat( os.path.join(args.dataset_folder, args.dataset_name, "clusters.mat")) clusters = mat['centroids'] import quantized_TF model = quantized_TF.QuantizedTF(clusters.shape[0], clusters.shape[0] + 1, clusters.shape[0], N=args.layers, d_model=args.emb_size, d_ff=1024, h=args.heads, dropout=args.dropout).to(device) tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) optim = NoamOpt( args.emb_size, args.factor, len(tr_dl) * 5, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)) #optim=Adagrad(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01,lr_decay=0.001) epoch = 0 while epoch < args.max_epoch: epoch_loss = 0 model.train() for id_b, batch in enumerate(tr_dl): optim.optimizer.zero_grad() scale = np.random.uniform(0.5, 4) #rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] * scale inp = torch.tensor( scipy.spatial.distance.cdist(speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape( n_in_batch, -1)).to(device) speeds_trg = batch['trg'][:, :, 2:4] * scale target = torch.tensor( scipy.spatial.distance.cdist(speeds_trg.reshape(-1, 2), clusters).argmin(axis=1).reshape( n_in_batch, -1)).to(device) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(target.shape[1]).repeat( n_in_batch, 1, 1).to(device) start_of_seq = torch.tensor( [clusters.shape[0]]).repeat(n_in_batch).unsqueeze(1).to(device) dec_inp = torch.cat((start_of_seq, target[:, :-1]), 1) out = model(inp, dec_inp, src_att, trg_att) loss = F.cross_entropy(out.view(-1, out.shape[-1]), target.view(-1), reduction='mean') loss.backward() optim.step() print("epoch %03i/%03i frame %04i / %04i loss: %7.4f" % (epoch, args.max_epoch, id_b, len(tr_dl), loss.item())) epoch_loss += loss.item() #sched.step() log.add_scalar('Loss/train', epoch_loss / len(tr_dl), epoch) with torch.no_grad(): model.eval() gt = [] pr = [] val_loss = 0 step = 0 for batch in val_dl: # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] inp = torch.tensor( scipy.spatial.distance.cdist( speeds_inp.contiguous().reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to(device) speeds_trg = batch['trg'][:, :, 2:4] target = torch.tensor( scipy.spatial.distance.cdist( speeds_trg.contiguous().reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to(device) src_att = torch.ones( (inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(target.shape[1]).repeat( n_in_batch, 1, 1).to(device) start_of_seq = torch.tensor([ clusters.shape[0] ]).repeat(n_in_batch).unsqueeze(1).to(device) dec_inp = torch.cat((start_of_seq, target[:, :-1]), 1) out = model(inp, dec_inp, src_att, trg_att) loss = F.cross_entropy(out.contiguous().view( -1, out.shape[-1]), target.contiguous().view(-1), reduction='mean') print("val epoch %03i/%03i frame %04i / %04i loss: %7.4f" % (epoch, args.max_epoch, step, len(val_dl), loss.item())) val_loss += loss.item() step += 1 log.add_scalar('validation/loss', val_loss / len(val_dl), epoch) if args.evaluate: # DETERMINISTIC MODE model.eval() model.eval() gt = [] pr = [] inp_ = [] peds = [] frames = [] dt = [] for batch in test_dl: inp_.append(batch['src'][:, :, 0:2]) gt.append(batch['trg'][:, :, 0:2]) frames.append(batch['frames']) peds.append(batch['peds']) dt.append(batch['dataset']) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] gt_b = batch['trg'][:, :, 0:2] inp = torch.tensor( scipy.spatial.distance.cdist( speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to(device) src_att = torch.ones( (inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(target.shape[1]).repeat( n_in_batch, 1, 1).to(device) start_of_seq = torch.tensor([ clusters.shape[0] ]).repeat(n_in_batch).unsqueeze(1).to(device) dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask(dec_inp.shape[1]).repeat( n_in_batch, 1, 1).to(device) out = model(inp, dec_inp, src_att, trg_att) dec_inp = torch.cat( (dec_inp, out[:, -1:].argmax(dim=2)), 1) preds_tr_b = clusters[dec_inp[:, 1:].cpu().numpy()].cumsum( 1) + batch['src'][:, -1:, 0:2].cpu().numpy() pr.append(preds_tr_b) peds = np.concatenate(peds, 0) frames = np.concatenate(frames, 0) dt = np.concatenate(dt, 0) gt = np.concatenate(gt, 0) dt_names = test_dataset.data['dataset_name'] pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) log.add_scalar('eval/DET_mad', mad, epoch) log.add_scalar('eval/DET_fad', fad, epoch) scipy.io.savemat( f"output/QuantizedTF/{args.name}/{epoch:05d}.mat", { 'input': inp, 'gt': gt, 'pr': pr, 'peds': peds, 'frames': frames, 'dt': dt, 'dt_names': dt_names }) # MULTI MODALITY if False: num_samples = 20 model.eval() gt = [] pr_all = {} for sam in range(num_samples): pr_all[sam] = [] for batch in test_dl: # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] gt_b = batch['trg'][:, :, 0:2] gt.append(gt_b) inp = torch.tensor( scipy.spatial.distance.cdist( speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape( n_in_batch, -1)).to(device) src_att = torch.ones( (inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(target.shape[1]).repeat( n_in_batch, 1, 1).to(device) start_of_seq = torch.tensor([ clusters.shape[0] ]).repeat(n_in_batch).unsqueeze(1).to(device) for sam in range(num_samples): dec_inp = start_of_seq for i in range(args.preds): trg_att = subsequent_mask( dec_inp.shape[1]).repeat(n_in_batch, 1, 1).to(device) out = model.predict(inp, dec_inp, src_att, trg_att) h = out[:, -1] dec_inp = torch.cat( (dec_inp, torch.multinomial(h, 1)), 1) preds_tr_b = clusters[dec_inp[:, 1:].cpu().numpy( )].cumsum(1) + batch['src'][:, -1:, 0:2].cpu().numpy() pr_all[sam].append(preds_tr_b) gt = np.concatenate(gt, 0) #pr=np.concatenate(pr,0) samp = {} for k in pr_all.keys(): samp[k] = {} samp[k]['pr'] = np.concatenate(pr_all[k], 0) samp[k]['mad'], samp[k]['fad'], samp[k][ 'err'] = baselineUtils.distance_metrics( gt, samp[k]['pr']) ev = [samp[i]['err'] for i in range(num_samples)] e20 = np.stack(ev, -1) mad_samp = e20.mean(1).min(-1).mean() fad_samp = e20[:, -1].min(-1).mean() #mad,fad,errs=baselineUtils.distance_metrics(gt,pr) log.add_scalar('eval/MM_mad', mad_samp, epoch) log.add_scalar('eval/MM_fad', fad_samp, epoch) if epoch % args.save_step == 0: torch.save(model.state_dict(), f'models/QuantizedTF/{args.name}/{epoch:05d}.pth') epoch += 1 ab = 1
def main(): parser = argparse.ArgumentParser( description='Train the individual Transformer model') parser.add_argument('--dataset_folder', type=str, default='datasets') parser.add_argument('--dataset_name', type=str, default='eth') parser.add_argument('--obs', type=int, default=8) parser.add_argument('--preds', type=int, default=12) parser.add_argument('--emb_size', type=int, default=1024) parser.add_argument('--heads', type=int, default=8) parser.add_argument('--layers', type=int, default=6) parser.add_argument('--dropout', type=float, default=0.1) parser.add_argument('--cpu', action='store_true') parser.add_argument('--output_folder', type=str, default='Output') parser.add_argument('--val_size', type=int, default=50) parser.add_argument('--gpu_device', type=str, default="0") parser.add_argument('--verbose', action='store_true') parser.add_argument('--max_epoch', type=int, default=100) parser.add_argument('--batch_size', type=int, default=256) parser.add_argument('--validation_epoch_start', type=int, default=30) parser.add_argument('--resume_train', action='store_true') parser.add_argument('--delim', type=str, default='\t') parser.add_argument('--name', type=str, default="eth_0.1") parser.add_argument('--factor', type=float, default=0.1) parser.add_argument('--save_step', type=int, default=1) args = parser.parse_args() model_name = args.name try: os.mkdir('models') except: pass try: os.mkdir('output') except: pass try: os.mkdir('output/BERT_quantized') except: pass try: os.mkdir(f'models/BERT_quantized') except: pass try: os.mkdir(f'output/BERT_quantized/{args.name}') except: pass try: os.mkdir(f'models/BERT_quantized/{args.name}') except: pass log = SummaryWriter('logs/BERT_quant_%s' % model_name) log.add_scalar('eval/mad', 0, 0) log.add_scalar('eval/fad', 0, 0) try: os.mkdir(args.name) except: pass device = torch.device("cuda") if args.cpu or not torch.cuda.is_available(): device = torch.device("cpu") args.verbose = True ## creation of the dataloaders for train and validation train_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=True, verbose=args.verbose) val_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, verbose=args.verbose) test_dataset, _ = baselineUtils.create_dataset(args.dataset_folder, args.dataset_name, 0, args.obs, args.preds, delim=args.delim, train=False, eval=True, verbose=args.verbose) #model.set_output_embeddings(GeneratorTS(1024,2)) tr_dl = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) val_dl = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0) test_dl = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) #optim = SGD(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01) #sched=torch.optim.lr_scheduler.StepLR(optim,0.0005) #optim=Adagrad(list(a.parameters())+list(model.parameters())+list(generator.parameters()),lr=0.01,lr_decay=0.001) epoch = 0 mat = scipy.io.loadmat( os.path.join(args.dataset_folder, args.dataset_name, "clusters.mat")) clusters = mat['centroids'] config = transformers.BertConfig(vocab_size=clusters.shape[0] + 1) gen = nn.Linear(config.hidden_size, clusters.shape[0]).to(device) model = transformers.BertModel(config).to(device) gen = nn.Linear(config.hidden_size, clusters.shape[0]).to(device) optim = NoamOpt( args.emb_size, args.factor, len(tr_dl) * 5, torch.optim.Adam(list(model.parameters()) + list(gen.parameters()), lr=0, betas=(0.9, 0.98), eps=1e-9)) mean = train_dataset[:]['src'][:, :, 2:4].mean((0, 1)) * 0 std = train_dataset[:]['src'][:, :, 2:4].std((0, 1)) * 0 + 1 while epoch < args.max_epoch: epoch_loss = 0 model.train() for id_b, batch in enumerate(tr_dl): optim.optimizer.zero_grad() scale = np.random.uniform(0.5, 2) # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] * scale inp = torch.tensor( scipy.spatial.distance.cdist(speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape( n_in_batch, -1)).to(device) speeds_trg = batch['trg'][:, :, 2:4] * scale target = torch.tensor( scipy.spatial.distance.cdist(speeds_trg.reshape(-1, 2), clusters).argmin(axis=1).reshape( n_in_batch, -1)).to(device) src_att = torch.ones((inp.shape[0], 1, inp.shape[1])).to(device) trg_att = subsequent_mask(target.shape[1]).repeat( n_in_batch, 1, 1).to(device) dec_inp = torch.tensor([clusters.shape[0] ]).repeat(n_in_batch, args.preds).to(device) bert_inp = torch.cat((inp, dec_inp), 1) out = gen( model(bert_inp, attention_mask=torch.ones( bert_inp.shape[0], bert_inp.shape[1]).to(device))[0]) loss = F.cross_entropy(out.view(-1, out.shape[-1]), torch.cat((inp, target), 1).view(-1), reduction='mean') loss.backward() optim.step() print("epoch %03i/%03i frame %04i / %04i loss: %7.4f" % (epoch, args.max_epoch, id_b, len(tr_dl), loss.item())) epoch_loss += loss.item() #sched.step() log.add_scalar('Loss/train', epoch_loss / len(tr_dl), epoch) with torch.no_grad(): model.eval() gt = [] pr = [] for batch in val_dl: gt_b = batch['trg'][:, :, 0:2] optim.optimizer.zero_grad() # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] inp = torch.tensor( scipy.spatial.distance.cdist( speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to(device) dec_inp = torch.tensor([clusters.shape[0] ]).repeat(n_in_batch, args.preds).to(device) bert_inp = torch.cat((inp, dec_inp), 1) out = gen( model(bert_inp, attention_mask=torch.ones( bert_inp.shape[0], bert_inp.shape[1]).to(device))[0]) F.softmax(out) preds_tr_b = clusters[F.softmax(out, dim=-1).argmax( dim=-1).cpu().numpy()][:, -args.preds:].cumsum( axis=1) + batch['src'][:, -1:, 0:2].cpu().numpy() gt.append(gt_b) pr.append(preds_tr_b) gt = np.concatenate(gt, 0) pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) log.add_scalar('validation/mad', mad, epoch) log.add_scalar('validation/fad', fad, epoch) model.eval() gt = [] pr = [] for batch in test_dl: gt_b = batch['trg'][:, :, 0:2] optim.optimizer.zero_grad() # rot_mat = np.array([[np.cos(r), np.sin(r)], [-np.sin(r), np.cos(r)]]) n_in_batch = batch['src'].shape[0] speeds_inp = batch['src'][:, 1:, 2:4] inp = torch.tensor( scipy.spatial.distance.cdist( speeds_inp.reshape(-1, 2), clusters).argmin(axis=1).reshape(n_in_batch, -1)).to(device) dec_inp = torch.tensor([clusters.shape[0] ]).repeat(n_in_batch, args.preds).to(device) bert_inp = torch.cat((inp, dec_inp), 1) out = gen( model(bert_inp, attention_mask=torch.ones( bert_inp.shape[0], bert_inp.shape[1]).to(device))[0]) F.softmax(out) preds_tr_b = clusters[F.softmax(out, dim=-1).argmax( dim=-1).cpu().numpy()][:, -args.preds:].cumsum( axis=1) + batch['src'][:, -1:, 0:2].cpu().numpy() gt.append(gt_b) pr.append(preds_tr_b) gt = np.concatenate(gt, 0) pr = np.concatenate(pr, 0) mad, fad, errs = baselineUtils.distance_metrics(gt, pr) if epoch % args.save_step == 0: torch.save( model.state_dict(), "models/BERT_quantized/%s/model_%03i.pth" % (args.name, epoch)) torch.save( gen.state_dict(), "models/BERT_quantized/%s/gen_%03i.pth" % (args.name, epoch)) log.add_scalar('eval/DET_mad', mad, epoch) log.add_scalar('eval/DET_fad', fad, epoch) epoch += 1 ab = 1