Speed up saving checkpoints
Summary:
Pull Request resolved:
It's better to write one checkpoint and copy it, rather than repeatedly pickling the model via torch.save.
@@ -14,6 +14,7 @@ import itertools
import math
import os
import random
+import shutil
import torch
@@ -307,8 +308,9 @@ def save_checkpoint(args, trainer, epoch_itr, val_loss):
checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
if len(checkpoints) > 0:
- for cp in checkpoints:
- trainer.save_checkpoint(cp, extra_state)
+ trainer.save_checkpoint(checkpoints[0], extra_state)
+ for cp in checkpoints[1:]:
+ shutil.copyfile(checkpoints[0], cp)
write_timer.stop()
print('| saved checkpoint {} (epoch {} @ {} updates) (writing took {} seconds)'.format(
