wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

added safeguards for callbacks in loading saving

Browse Source
This commit is contained in:
William Falcon
2019-07-24 11:14:19 -04:00
parent a4b8aa0a41
commit 8fd7a6001b
1 changed files with 15 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/root_module/model_saving.py
+15 -6
View File
@@ -51,14 +51,19 @@ class TrainerIO(object):
torch.save(checkpoint, filepath)
def dump_checkpoint(self):
checkpoint = {
'epoch': self.current_epoch,
'checkpoint_callback_best': self.checkpoint_callback.best,
'early_stop_callback_wait': self.early_stop_callback.wait,
'early_stop_callback_patience': self.early_stop_callback.patience,
'global_step': self.global_step
}
if self.checkpoint_callback is not None:
checkpoint['checkpoint_callback_best'] = self.checkpoint_callback_best.best
if self.early_stop_callback is not None:
checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
optimizer_states = []
for i, optimizer in enumerate(self.optimizers):
optimizer_states.append(optimizer.state_dict())
@@ -104,9 +109,13 @@ class TrainerIO(object):
:param checkpoint:
:return:
"""
self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
if self.checkpoint_callback is not None:
self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
if self.early_stop_callback is not None:
self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
self.global_step = checkpoint['global_step']
self.current_epoch = checkpoint['epoch']