past checkpoints (Lightning-AI#2160)

* past checkpoints

* omegaConf save

* enforce type

* resolve=True

Co-authored-by: Omry Yadan <omry@fb.com>

* test omegaconf

* tests

* test past

Co-authored-by: Omry Yadan <omry@fb.com>

Author: Borda

pytorch_lightning/core/saving.py

@@ -21,10 +21,17 @@
 else:
     ALLOWED_CONFIG_TYPES = ALLOWED_CONFIG_TYPES + (Container, )
 
+# the older shall be on the top
+CHECKPOINT_PAST_HPARAMS_KEYS = (
+    'hparams',
+    'module_arguments',  # used in 0.7.6
+)
+
 
 class ModelIO(object):
-    CHECKPOINT_KEY_HYPER_PARAMS = 'hyper_parameters'
-    CHECKPOINT_NAME_HYPER_PARAMS = 'hparams_name'
+    CHECKPOINT_HYPER_PARAMS_KEY = 'hyper_parameters'
+    CHECKPOINT_HYPER_PARAMS_NAME = 'hparams_name'
+    CHECKPOINT_HYPER_PARAMS_TYPE = 'hparams_type'
 
     @classmethod
     def load_from_metrics(cls, weights_path, tags_csv, map_location=None):
@@ -153,21 +160,29 @@ def load_from_checkpoint(
             hparams['on_gpu'] = False
 
             # overwrite hparams by the given file
-            checkpoint[cls.CHECKPOINT_KEY_HYPER_PARAMS] = hparams
+            checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
 
-        # override the module_arguments with values that were passed in
-        checkpoint[cls.CHECKPOINT_KEY_HYPER_PARAMS].update(kwargs)
+        # for past checkpoint need to add the new key
+        if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
+            checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
+        # override the hparams with values that were passed in
+        checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
 
         model = cls._load_model_state(checkpoint, *args, **kwargs)
         return model
 
     @classmethod
     def _load_model_state(cls, checkpoint: Dict[str, Any], *args, **kwargs):
         # pass in the values we saved automatically
-        if cls.CHECKPOINT_KEY_HYPER_PARAMS in checkpoint:
-            # todo add some back compatibility
-            model_args = checkpoint[cls.CHECKPOINT_KEY_HYPER_PARAMS]
-            args_name = checkpoint.get(cls.CHECKPOINT_NAME_HYPER_PARAMS)
+        if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
+            model_args = {}
+            # add some back compatibility, the actual one shall be last
+            for hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS + (cls.CHECKPOINT_HYPER_PARAMS_KEY,):
+                if hparam_key in checkpoint:
+                    model_args.update(checkpoint[hparam_key])
+            if cls.CHECKPOINT_HYPER_PARAMS_TYPE in checkpoint:
+                model_args = checkpoint[cls.CHECKPOINT_HYPER_PARAMS_TYPE](model_args)
+            args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)
             init_args_name = inspect.signature(cls).parameters.keys()
             if args_name == 'kwargs':
                 cls_kwargs = {k: v for k, v in model_args.items() if k in init_args_name}

pytorch_lightning/loggers/tensorboard.py

@@ -4,7 +4,6 @@
 """
 
 import os
-import yaml
 from argparse import Namespace
 from typing import Optional, Dict, Union, Any
 from warnings import warn
@@ -18,6 +17,11 @@
 from pytorch_lightning.loggers.base import LightningLoggerBase
 from pytorch_lightning.utilities import rank_zero_only
 
+try:
+    from omegaconf import Container
+except ImportError:
+    Container = None
+
 
 class TensorBoardLogger(LightningLoggerBase):
     r"""
@@ -152,7 +156,14 @@ def save(self) -> None:
         hparams_file = os.path.join(dir_path, self.NAME_HPARAMS_FILE)
 
         # save the metatags file
-        save_hparams_to_yaml(hparams_file, self.hparams)
+        if Container is not None:
+            if isinstance(self.hparams, Container):
+                from omegaconf import OmegaConf
+                OmegaConf.save(self.hparams, hparams_file, resolve=True)
+            else:
+                save_hparams_to_yaml(hparams_file, self.hparams)
+        else:
+            save_hparams_to_yaml(hparams_file, self.hparams)
 
     @rank_zero_only
     def finalize(self, status: str) -> None:

pytorch_lightning/trainer/training_io.py

@@ -87,20 +87,20 @@
 import re
 import signal
 from abc import ABC
-from argparse import Namespace
 from subprocess import call
 
 import torch
 import torch.distributed as torch_distrib
 
+import pytorch_lightning
 from pytorch_lightning import _logger as log
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.loggers import LightningLoggerBase
 from pytorch_lightning.overrides.data_parallel import (
     LightningDistributedDataParallel,
     LightningDataParallel,
 )
-from pytorch_lightning.utilities import rank_zero_warn, parsing
+from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.io import load as pl_load
 
 try:
@@ -119,6 +119,11 @@
 else:
     HOROVOD_AVAILABLE = True
 
+try:
+    from omegaconf import Container
+except ImportError:
+    Container = None
+
 
 class TrainerIOMixin(ABC):
 
@@ -267,8 +272,8 @@ def save_checkpoint(self, filepath, weights_only: bool = False):
             try:
                 self._atomic_save(checkpoint, filepath)
             except AttributeError as err:
-                if LightningModule.CHECKPOINT_KEY_HYPER_PARAMS in checkpoint:
-                    del checkpoint[LightningModule.CHECKPOINT_KEY_HYPER_PARAMS]
+                if LightningModule.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
+                    del checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]
                 rank_zero_warn('Warning, `module_arguments` dropped from checkpoint.'
                                f' An attribute is not picklable {err}')
                 self._atomic_save(checkpoint, filepath)
@@ -320,6 +325,7 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
         checkpoint = {
             'epoch': self.current_epoch + 1,
             'global_step': self.global_step + 1,
+            'pytorch-ligthning_version': pytorch_lightning.__version__,
         }
 
         if not weights_only:
@@ -356,10 +362,12 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
 
         if model.hparams:
             if hasattr(model, '_hparams_name'):
-                checkpoint[LightningModule.CHECKPOINT_NAME_HYPER_PARAMS] = model._hparams_name
+                checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_NAME] = model._hparams_name
             # add arguments to the checkpoint
-            # todo: add some recursion in case of OmegaConf
-            checkpoint[LightningModule.CHECKPOINT_KEY_HYPER_PARAMS] = dict(model.hparams)
+            checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY] = model.hparams
+            if Container is not None:
+                if isinstance(model.hparams, Container):
+                    checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_TYPE] = type(model.hparams)
 
         # give the model a chance to add a few things
         model.on_save_checkpoint(checkpoint)
@@ -473,8 +481,8 @@ def hpc_save(self, folderpath: str, logger):
         try:
             self._atomic_save(checkpoint, filepath)
         except AttributeError as err:
-            if LightningModule.CHECKPOINT_KEY_HYPER_PARAMS in checkpoint:
-                del checkpoint[LightningModule.CHECKPOINT_KEY_HYPER_PARAMS]
+            if LightningModule.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
+                del checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]
             rank_zero_warn('warning, `module_arguments` dropped from checkpoint.'
                            f' An attribute is not picklable {err}')
             self._atomic_save(checkpoint, filepath)