Add resuming from specific checkpoint

pytorch_lightning/trainer/trainer.py

@@ -82,7 +82,8 @@ def __init__(self,
  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
@@ -119,6 +120,7 @@ def __init__(self,
  :param amp_level: str. Check nvidia docs for level
  :param nb_sanity_val_steps: int. How many val steps before a full train loop.
  :param truncated_bptt_steps: int. Enables multiple backward passes for each batch.
+ :param resume_from_checkpoint: str or os.PathLike object. Resume from specific checkpoint.
  """
  # Transfer params
  self.nb_gpu_nodes = nb_gpu_nodes
@@ -139,6 +141,7 @@ def __init__(self,
  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

pytorch_lightning/trainer/trainer_io.py

@@ -2,6 +2,7 @@
 import re
 import signal
 import warnings
+from pathlib import Path
 from subprocess import call
 import logging
 
@@ -46,7 +47,9 @@ def restore_weights(self, model):
 
  if not did_restore_hpc_weights:
  # restore weights if same exp version
- self.restore_state_if_checkpoint_exists(model)
+ did_restore_last_checkpoint = self.restore_state_if_checkpoint_exists(model)
+ if not did_restore_last_checkpoint and self.resume_from_checkpoint is not None:
+ self.restore_state_from_checkpoint(self.resume_from_checkpoint)
 
  # wait for all models to restore weights
  if self.use_ddp or self.use_ddp2:
@@ -93,6 +96,18 @@ def restore_state_if_checkpoint_exists(self, model):
 
  return did_restore
 
+ def restore_state_from_checkpoint(self, checkpoint_path):
+ did_restore = False
+
+ checkpoint_path = Path(checkpoint_path)
+ if not checkpoint_path.exists():
+ return did_restore
+
+ self.restore(checkpoint_path, self.on_gpu)
+ did_restore = True
+
+ return did_restore
+
  # --------------------
  # HPC SIGNAL HANDLING
  # --------------------