Repository: Lightning-AI/pytorch-lightning

Issue:
Customizing hparams after loading checkpoint
## ❓ Questions and Help

### Before asking:   
1. search the issues.   
2. search the docs.    

#### What is your question?
I'm wondering what the best practice for loading a model with different hparams than what is stored in the checkpoint?

I realize I could just load the model and set them afterwards e.g.: 
```
model = model.load_from_checkpoint(args.checkpoint_file) # Load model

# Set hparams etc..
model.hparams.arg1 = 0.0
model.hparams.arg2 = 1.0
 
```
But the problem is that my  model __init__ function depends on the hparams arg1 and arg2 so they're set too late.

I could also do 
```
checkpoint = torch.load(args.checkpoint_file)
checkpoint['hparams']['arg1'] = 0.0
checkpoint['hparams']['arg2'] = 1.0
model = model._load_state_dict(checkpoint)
```

The problem here is that i'm using the protected function _load_state_dict. Is there another way of solving this that i've missed? Or could we consider making _load_state_dict public?



Discussion:
@tullie good question, maybe we can have a flag for disabling hparam use?
```python
load_from_checkpoint(PATH, auto_hparam=False, arg1=my_arg, arg2=my_arg2, hparam={...})
```

load_from_checkpoint currently allows passing in the args directly as shown above
```python
load_from_checkpoint(PATH, auto_hparam=False, arg1=my_arg, arg2=my_arg2)
```
so, in this case an arg (hparam) would be a regular arg which you could then construct to be whatever you want.

Alternative 2:
We add a ```hparam_updates``` arg which sets those updates in the hparams

```python
load_from_checkpoint(PATH, hparam_updates={'arg1': 0.0, 'arg2': 0.0})
```
I've been playing around with both these options and have most preferred alternative 2.
I'll send out a PR soon. 

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -16,8 +16,8 @@
 from pytorch_lightning.core.grads import GradInformation
 from pytorch_lightning.core.hooks import ModelHooks
 from pytorch_lightning.core.memory import ModelSummary
+from pytorch_lightning.core.saving import ModelIO, load_hparams_from_tags_csv, update_hparams
 from pytorch_lightning.core.properties import DeviceDtypeModuleMixin
-from pytorch_lightning.core.saving import ModelIO, load_hparams_from_tags_csv
 from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities import rank_zero_warn
@@ -1439,6 +1439,7 @@ def load_from_checkpoint(
             checkpoint_path: str,
             map_location: Optional[Union[Dict[str, str], str, torch.device, int, Callable]] = None,
             tags_csv: Optional[str] = None,
+            hparam_overrides: Optional[Dict] = None,
             *args, **kwargs
     ) -> 'LightningModule':
         r"""
@@ -1480,6 +1481,7 @@ def __init__(self, hparams):
                 use this method to pass in a .csv file with the hparams you'd like to use.
                 These will be converted into a :class:`~argparse.Namespace` and passed into your
                 :class:`LightningModule` for use.
+            hparam_overrides: A dictionary with keys to override in the hparams
 
         Return:
             :class:`LightningModule` with loaded weights and hyperparameters (if available).
@@ -1503,6 +1505,12 @@ def __init__(self, hparams):
                     tags_csv='/path/to/hparams_file.csv'
                 )
 
+                # override some of the params with new values
+                MyLightningModule.load_from_checkpoint(
+                    PATH,
+                    hparam_overrides={'num_layers': 128, 'pretrained_ckpt_path': NEW_PATH}
+                )
+
                 # or load passing whatever args the model takes to load
                 MyLightningModule.load_from_checkpoint(
                     'path/to/checkpoint.ckpt',
@@ -1521,12 +1529,16 @@ def __init__(self, hparams):
         else:
             checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
 
+        # add the hparams from csv file to checkpoint
         if tags_csv is not None:
-            # add the hparams from csv file to checkpoint
             hparams = load_hparams_from_tags_csv(tags_csv)
             hparams.__setattr__('on_gpu', False)
             checkpoint['hparams'] = vars(hparams)
 
+        # override the hparam keys that were passed in
+        if hparam_overrides is not None:
+            update_hparams(hparams, hparam_overrides)
+
         model = cls._load_model_state(checkpoint, *args, **kwargs)
         return model
 
diff --git a/pytorch_lightning/core/saving.py b/pytorch_lightning/core/saving.py
--- a/pytorch_lightning/core/saving.py
+++ b/pytorch_lightning/core/saving.py
@@ -48,6 +48,37 @@ def on_hpc_load(self, checkpoint: Dict[str, Any]) -> None:
         """
 
 
+def update_hparams(hparams: dict, updates: dict) -> None:
+    """
+    Overrides hparams with new values
+
+    >>> hparams = {'c': 4}
+    >>> update_hparams(hparams, {'a': {'b': 2}, 'c': 1})
+    >>> hparams['a']['b'], hparams['c']
+    (2, 1)
+    >>> update_hparams(hparams, {'a': {'b': 4}, 'c': 7})
+    >>> hparams['a']['b'], hparams['c']
+    (4, 7)
+
+    Args:
+        hparams: the original params and also target object
+        updates: new params to be used as update
+
+    """
+    for k, v in updates.items():
+        # if missing, add the key
+        if k not in hparams:
+            hparams[k] = v
+            continue
+
+        # recurse if dictionary
+        if isinstance(v, dict):
+            update_hparams(hparams[k], updates[k])
+        else:
+            # update the value
+            hparams.update({k: v})
+
+
 def load_hparams_from_tags_csv(tags_csv: str) -> Namespace:
     if not os.path.isfile(tags_csv):
         log.warning(f'Missing Tags: {tags_csv}.')


Test patch: