Aquila-Consortium/map2map
32
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Save past best models by not overwriting them

Browse Source
This commit is contained in:
Yin Li 2019-12-12 16:51:47 -05:00
parent 946805c6be
commit 0533150194
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
map2map/train.py
View File

@ -144,13 +144,16 @@ def gpu_worker(local_rank, args):
'rng' : torch.get_rng_state(), 'rng' : torch.get_rng_state(),
'min_loss': min_loss, 'min_loss': min_loss,
} }
filename='checkpoint.pth' ckpt_file = 'checkpoint.pth'
best_file = 'best_model_{}.pth'
torch.save(state, filename) torch.save(state, filename)
del state del state
if min_loss is None or val_loss < min_loss: if min_loss is None or val_loss < min_loss:
min_loss = val_loss min_loss = val_loss
shutil.copyfile(filename, 'best_model.pth') shutil.copyfile(filename, best_file.format(epoch + 1))
if os.path.isfile(best_file.format(epoch)):
os.remove(best_file.format(epoch))
destroy_process_group() destroy_process_group()