Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added atomic checkpoint creation #689

Merged
merged 2 commits into from
Jan 20, 2020
Merged

Added atomic checkpoint creation #689

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b52365a
Added atomic checkpoint creation
Jan 15, 2020
eda52d5
Added documentation for _atomic_checkpoint
Jan 20, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions pytorch_lightning/trainer/training_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,17 +255,22 @@ def term_handler(self, signum, frame):
# --------------------
# MODEL SAVE CHECKPOINT
# --------------------
def _atomic_save(self, checkpoint, filepath):
fgerzer marked this conversation as resolved.
Show resolved Hide resolved
fgerzer marked this conversation as resolved.
Show resolved Hide resolved
tmp_path = str(filepath) + ".part"
torch.save(checkpoint, tmp_path)
os.replace(tmp_path, filepath)

def save_checkpoint(self, filepath):
checkpoint = self.dump_checkpoint()

# do the actual save
try:
torch.save(checkpoint, filepath)
self._atomic_save(checkpoint, filepath)
except AttributeError:
if 'hparams' in checkpoint:
del checkpoint['hparams']

torch.save(checkpoint, filepath)
self._atomic_save(checkpoint, filepath)

def restore(self, checkpoint_path, on_gpu):
# if on_gpu:
Expand Down Expand Up @@ -412,12 +417,12 @@ def hpc_save(self, folderpath, logger):
# do the actual save
# TODO: fix for anything with multiprocess DP, DDP, DDP2
try:
torch.save(checkpoint, filepath)
self._atomic_save(checkpoint, filepath)
except AttributeError:
if 'hparams' in checkpoint:
del checkpoint['hparams']

torch.save(checkpoint, filepath)
self._atomic_save(checkpoint, filepath)

return filepath

Expand Down