From 2b8475f5909dff3d9b0c83c7b00aa15bfda63ee1 Mon Sep 17 00:00:00 2001 From: Yongrae Jo Date: Sun, 1 Dec 2019 06:48:38 +0900 Subject: [PATCH] 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 --- pytorch_lightning/trainer/trainer.py | 4 +++- pytorch_lightning/trainer/trainer_io.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) 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: