diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index fc340859..d3fa337f 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -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 diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index 7e5ce7e4..4bd32e81 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -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: