Rooba
/
genienlp
mirror of https://github.com/stanford-oval/genienlp.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

train: don't save per-worker checkpoints if we're not doing distributed training 

Saves disk space
This commit is contained in:
Giovanni Campagna 2019-03-01 10:51:12 -08:00
parent a36f2efb8c
commit d4b35d7ae6
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
train.py
View File

@ -215,8 +215,7 @@ def train(args, model, opt, train_iters, train_iterations, field, rank=0, world_
torch.save({'model_state_dict': {k: v.cpu() for k, v in model.state_dict().items()}, 'field': field}, os.path.join(args.log_dir, f'iteration_{iteration}.pth'))
if world_size > 1:
torch.distributed.barrier()
torch.save(opt.state_dict(), os.path.join(args.log_dir, f'iteration_{iteration}_rank_{rank}_optim.pth'))
if world_size > 1:
torch.save(opt.state_dict(), os.path.join(args.log_dir, f'iteration_{iteration}_rank_{rank}_optim.pth'))
torch.distributed.barrier()
# lr update