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
