Checkpoint Loading process

This includes the basics of the DDP setup

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}')  

Checkpoint Saving snippet

if (epoch + 1) % args.val_per_n_epoch == 0:
            if rank == 0:
                # save model
                ckpt_path = '/fsx/awesome/comma2k19_checkpoints/%s_%d.pth' % (args.model, epoch)
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.module.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'lr_scheduler_state_dict': lr_scheduler.state_dict(),
                    'loss': total_loss,
                }, ckpt_path)

                time.sleep(10) # wait for all processes to save the file
                wandb.save(ckpt_path)
                print('[Epoch %d] checkpoint saved at %s' % (epoch, ckpt_path))
Edit
Pub: 23 Oct 2022 22:30 UTC
Views: 66