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:
Yongrae Jo
2019-11-30 16:48:38 -05:00
committed by William Falcon
parent df7b6d958e
commit 2b8475f590
2 changed files with 8 additions and 3 deletions
+3 -1
View File
@@ -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
+5 -2
View File
@@ -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: