mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
Add resuming from specific checkpoint (#516)
* Add resume_from_checkpoint * Fix variable name * #515 Remove did_restore * #515 Simplify code * #515 Update doc for resume_from_checkpoint * #515 Add on_gpu
This commit is contained in:
committed by
William Falcon
parent
df7b6d958e
commit
2b8475f590
@@ -83,7 +83,8 @@ class Trainer(TrainerIOMixin,
|
||||
weights_save_path=None,
|
||||
amp_level='O1',
|
||||
nb_sanity_val_steps=5,
|
||||
truncated_bptt_steps=None):
|
||||
truncated_bptt_steps=None,
|
||||
resume_from_checkpoint=None):
|
||||
"""
|
||||
|
||||
:param logger: Logger for experiment tracking
|
||||
@@ -140,6 +141,7 @@ class Trainer(TrainerIOMixin,
|
||||
self.nb_sanity_val_steps = nb_sanity_val_steps
|
||||
self.print_nan_grads = print_nan_grads
|
||||
self.truncated_bptt_steps = truncated_bptt_steps
|
||||
self.resume_from_checkpoint = resume_from_checkpoint
|
||||
self.shown_warnings = set()
|
||||
|
||||
self.fast_dev_run = fast_dev_run
|
||||
|
||||
@@ -138,8 +138,11 @@ class TrainerIOMixin(object):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if not did_restore_hpc_weights:
|
||||
# restore weights if same exp version
|
||||
self.restore_state_if_checkpoint_exists(model)
|
||||
if self.resume_from_checkpoint is not None:
|
||||
self.restore(self.resume_from_checkpoint, on_gpu=self.on_gpu)
|
||||
else:
|
||||
# restore weights if same exp version
|
||||
self.restore_state_if_checkpoint_exists(model)
|
||||
|
||||
# wait for all models to restore weights
|
||||
if self.use_ddp or self.use_ddp2:
|
||||
|
||||
Reference in New Issue
Block a user