model = SequenceBaselineV1(args.model, args.M, args.num_pts, args.mtp_alpha, args.lr, args.optimizer, args.optimize_per_n_step)
use_sync_bn = args.sync_bn
if use_sync_bn:
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = model.cuda()
optimizer, lr_scheduler = model.configure_optimizers(args, model)
model: SequenceBaselineV1
dist.barrier()
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank], find_unused_parameters=True, broadcast_buffers=False)
loss = MultipleTrajectoryPredictionLoss(args.mtp_alpha, args.M, args.num_pts, distance_type='angle')
if args.resume:
# resuming mechanism, obtaining the latest checkpoint path
chkp_list = glob(f'/fsx/awesome/comma2k19_checkpoints/{args.model}*.pth')
chkp_file = max(chkp_list, key=lambda x: int(x.split('_')[-1].split('.')[0]))
if len(chkp_list) > 0:
# Load everything
checkpoint = torch.load(chkp_file, map_location=f'cuda:{rank}')
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler_state_dict'])
print(f'\n{"==="*25}\nLoaded checkpoint from {chkp_file}\n{"==="*25}')