st_time = time.time()
hist, nbrs, mask, lat_enc, lon_enc, fut, op_mask, vehid, t, ds = data
if args['use_cuda']:
hist = hist.cuda()
nbrs = nbrs.cuda()
mask = mask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
fut = fut.cuda()
op_mask = op_mask.cuda()
fut_pred, weight_ts_center, weight_ts_nbr, weight_ha = net(
hist, nbrs, mask, lat_enc, lon_enc)
l = maskedMSE(fut_pred, fut,
op_mask) #maskedNLL(fut_pred, fut, op_mask)
# Backprop and update weights
optimizer.zero_grad()
l.backward()
a = torch.nn.utils.clip_grad_norm_(net.parameters(), 10)
optimizer.step()
#optimizer = tf.train.RMSPropOptimizer(learning_rate, decay).minimize(cost)
# Track average train loss and average train time:
batch_time = time.time() - st_time
avg_tr_loss += l.item() # sum mse for 100 batches
avg_tr_time += batch_time
if i % 100 == 99:
eta = avg_tr_time / 100 * (len(trSet) / batch_size - i
) # average time/batch * rest batches
def train_model():
args = parser.parse_args()
print("------------- {} -------------".format(args.name))
print("Batch size : {}".format(args.batch_size))
print("Learning rate : {}".format(args.learning_rate))
print("Use Planning Coupled: {}".format(args.use_planning))
print("Use Target Fusion: {}".format(args.use_fusion))
## Initialize network and optimizer
PiP = pipNet(args)
if args.use_cuda:
PiP = PiP.cuda()
optimizer = torch.optim.Adam(PiP.parameters(), lr=args.learning_rate)
crossEnt = torch.nn.BCELoss()
## Initialize the log folder
log_path = "./trained_models/{}/".format(args.name)
if not os.path.exists(log_path):
os.makedirs(log_path)
if args.tensorboard:
logger = SummaryWriter(log_path + 'train-pre{}-nll{}'.format(args.pretrain_epochs, args.train_epochs))
logger_val = SummaryWriter(log_path + 'validation-pre{}-nll{}'.format(args.pretrain_epochs, args.train_epochs))
## Initialize training parameters
pretrainEpochs = args.pretrain_epochs
trainEpochs = args.train_epochs
batch_size = args.batch_size
## Initialize data loaders
print("Train dataset: {}".format(args.train_set))
trSet = highwayTrajDataset(path=args.train_set,
targ_enc_size=args.social_context_size+args.dynamics_encoding_size,
grid_size=args.grid_size,
fit_plan_traj=False)
print("Validation dataset: {}".format(args.val_set))
valSet = highwayTrajDataset(path=args.val_set,
targ_enc_size=args.social_context_size+args.dynamics_encoding_size,
grid_size=args.grid_size,
fit_plan_traj=True)
trDataloader = DataLoader(trSet, batch_size=batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=trSet.collate_fn)
valDataloader = DataLoader(valSet, batch_size=batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=valSet.collate_fn)
print("DataSet Prepared : {} train data, {} validation data\n".format(len(trSet), len(valSet)))
print("Network structure: {}\n".format(PiP))
## Training process
for epoch_num in range( pretrainEpochs + trainEpochs ):
if epoch_num == 0:
print('Pretrain with MSE loss')
elif epoch_num == pretrainEpochs:
print('Train with NLL loss')
## Variables to track training performance:
avg_time_tr, avg_loss_tr, avg_loss_val = 0, 0, 0
## Training status, reclaim after each epoch
PiP.train()
PiP.train_output_flag = True
for i, data in enumerate(trDataloader):
st_time = time.time()
nbsHist, nbsMask, planFut, planMask, targsHist, targsEncMask, targsFut, targsFutMask, lat_enc, lon_enc, _ = data
if args.use_cuda:
nbsHist = nbsHist.cuda()
nbsMask = nbsMask.cuda()
planFut = planFut.cuda()
planMask = planMask.cuda()
targsHist = targsHist.cuda()
targsEncMask = targsEncMask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
targsFut = targsFut.cuda()
targsFutMask = targsFutMask.cuda()
# Forward pass
fut_pred, lat_pred, lon_pred = PiP(nbsHist, nbsMask, planFut, planMask, targsHist, targsEncMask, lat_enc, lon_enc)
if epoch_num < pretrainEpochs:
# Pre-train with MSE loss to speed up training
l = maskedMSE(fut_pred, targsFut, targsFutMask)
else:
# Train with NLL loss
l = maskedNLL(fut_pred, targsFut, targsFutMask) + crossEnt(lat_pred, lat_enc) + crossEnt(lon_pred, lon_enc)
# Back-prop and update weights
optimizer.zero_grad()
l.backward()
prev_vec_norm = torch.nn.utils.clip_grad_norm_(PiP.parameters(), 10)
optimizer.step()
# Track average train loss and average train time:
batch_time = time.time()-st_time
avg_loss_tr += l.item()
avg_time_tr += batch_time
# For every 100 batches: record loss, validate model, and plot.
if i%100 == 99:
eta = avg_time_tr/100*(len(trSet)/batch_size-i)
epoch_progress = i * batch_size / len(trSet)
print("Epoch no:",epoch_num+1,
"| Epoch progress(%):",format(epoch_progress*100,'0.2f'),
"| Avg train loss:",format(avg_loss_tr/100,'0.2f'),
"| ETA(s):",int(eta))
if args.tensorboard:
logger.add_scalar("RMSE" if epoch_num < pretrainEpochs else "NLL", avg_loss_tr / 100, (epoch_progress + epoch_num) * 100)
## Validatation during training:
eval_batch_num = 20
with torch.no_grad():
PiP.eval()
PiP.train_output_flag = False
for i, data in enumerate(valDataloader):
nbsHist, nbsMask, planFut, planMask, targsHist, targsEncMask, targsFut, targsFutMask, lat_enc, lon_enc, _ = data
if args.use_cuda:
nbsHist = nbsHist.cuda()
nbsMask = nbsMask.cuda()
planFut = planFut.cuda()
planMask = planMask.cuda()
targsHist = targsHist.cuda()
targsEncMask = targsEncMask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
targsFut = targsFut.cuda()
targsFutMask = targsFutMask.cuda()
if epoch_num < pretrainEpochs:
# During pre-training with MSE loss, validate with MSE for true maneuver class trajectory
PiP.train_output_flag = True
fut_pred, _, _ = PiP(nbsHist, nbsMask, planFut, planMask, targsHist, targsEncMask,
lat_enc, lon_enc)
l = maskedMSE(fut_pred, targsFut, targsFutMask)
else:
# During training with NLL loss, validate with NLL over multi-modal distribution
fut_pred, lat_pred, lon_pred = PiP(nbsHist, nbsMask, planFut, planMask, targsHist,
targsEncMask, lat_enc, lon_enc)
l = maskedNLLTest(fut_pred, lat_pred, lon_pred, targsFut, targsFutMask, avg_along_time=True)
avg_loss_val += l.item()
if i==(eval_batch_num-1):
if args.tensorboard:
logger_val.add_scalar("RMSE" if epoch_num < pretrainEpochs else "NLL", avg_loss_val / eval_batch_num, (epoch_progress + epoch_num) * 100)
break
# Clear statistic
avg_time_tr, avg_loss_tr, avg_loss_val = 0, 0, 0
# Revert to train mode after in-process evaluation.
PiP.train()
PiP.train_output_flag = True
## Save the model after each epoch______________________________________________________________________________
epoCount = epoch_num + 1
if epoCount < pretrainEpochs:
torch.save(PiP.state_dict(), log_path + "{}-pre{}-nll{}.tar".format(args.name, epoCount, 0))
else:
torch.save(PiP.state_dict(), log_path + "{}-pre{}-nll{}.tar".format(args.name, pretrainEpochs, epoCount - pretrainEpochs))
# All epochs finish________________________________________________________________________________________________
torch.save(PiP.state_dict(), log_path+"{}.tar".format(args.name))
print("Model saved in trained_models/{}/{}.tar\n".format(args.name, args.name))
hist = hist.cuda()
nbrs = nbrs.cuda()
mask = mask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
fut = fut.cuda()
op_mask = op_mask.cuda()
hist_grid = hist_grid.cuda()
# Forward pass
if params.use_maneuvers:
fut_pred, lat_pred, lon_pred = net(hist, nbrs, mask, lat_enc,
lon_enc, hist_grid, fut)
# Pre-train with MSE loss to speed up training
if epoch_num < pretrainEpochs:
l = maskedMSE(fut_pred, fut, op_mask)
else:
# Train with NLL loss
l = maskedNLL(fut_pred, fut, op_mask) + crossEnt(
lat_pred, lat_enc) + crossEnt(lon_pred, lon_enc)
avg_lat_acc += (torch.sum(
torch.max(lat_pred.data, 1)[1] == torch.max(
lat_enc.data, 1)[1])).item() / lat_enc.size()[0]
avg_lon_acc += (torch.sum(
torch.max(lon_pred.data, 1)[1] == torch.max(
lon_enc.data, 1)[1])).item() / lon_enc.size()[0]
else:
fut_pred = net(hist, nbrs, mask, lat_enc, lon_enc, hist_grid, fut)
if epoch_num < pretrainEpochs:
l = maskedMSE(fut_pred, fut, op_mask)
else:
tic = time.time()
for ep in range(args['train_epoches']):
running_loss = 0.0
for i, data in enumerate(trDataloader):
hh, ff = data
hh = hh[:, ::2, :].to(device)
ff = ff[:, 4::5, :].to(device)
## zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
fut_pred = net(hh)
op_mask = torch.ones(ff.shape)
l = maskedMSE(fut_pred, ff, op_mask)
l.backward()
# a = torch.nn.utils.clip_grad_norm_(net.parameters(), 10)
optimizer.step()
# print statistics
running_loss += l.item()
if i % 1000 == 999: # print every 1000 mini-batches
print('ep {}, {} batches, loss - {}'.format(
ep + 1, i + 1, running_loss / 1000))
running_loss = 0.0
###_____________________________________________________________________________________________________________________________________________
## no validation
## save model