Repository: Lightning-AI/pytorch-lightning

Issue:
Checkpoints cannot be loaded in non-pl env
~~## 🚀 Feature
Add an option to save only `state_dict` for `ModelCheckpoint` callbacks~~

## 🐛 Bug
PL checkpoints cannot be loaded in non-pl envs

### Motivation
To be able to move trained models and weights into pytorch only environments

### Additional context
Currently  when you do `torch.load()` on a pl generated checkpoint in an environment without pl, there is a pickling error. For my current use case I have to load the checkpoints in my training environment and save them again with only `state_dict` for the weights.

See [reply below](https://github.com/PyTorchLightning/pytorch-lightning/issues/2653#issuecomment-681303007) for more info


Discussion:
You can use `save_weights_only` parameter in ModelCheckpoint to save weights only. Although it will save `epoch`, `global_step` and `pl_version` but that won't be a problem there, I guess. Also can you show the pickling error you are getting?
I am using `save_weights_only ` and that causes a pickling error with module lightning not found (don't have the extact error atm)
Can you check when you load that checkpoint manually in pl env, what keys does that file have?
Error in non-pl env
```
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-10-dbc5018f5317> in <module>
----> 1 pretrained_dict = torch.load('../input/weights/test.ckpt', map_location=torch.device('cpu'))

/opt/conda/lib/python3.7/site-packages/torch/serialization.py in load(f, map_location, pickle_module, **pickle_load_args)
    591                     return torch.jit.load(f)
    592                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
--> 593         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
    594 
    595 

/opt/conda/lib/python3.7/site-packages/torch/serialization.py in _legacy_load(f, map_location, pickle_module, **pickle_load_args)
    771     unpickler = pickle_module.Unpickler(f, **pickle_load_args)
    772     unpickler.persistent_load = persistent_load
--> 773     result = unpickler.load()
    774 
    775     deserialized_storage_keys = pickle_module.load(f, **pickle_load_args)

ModuleNotFoundError: No module named 'pytorch_lightning'
```

keys in pl env
`dict_keys(['epoch', 'global_step', 'pytorch-lightning_version', 'state_dict', 'hparams_name', 'hyper_parameters'])`

@rohitgr7 sorry about the late reply, completely forgot about this issue

Edit:
found the issue, I'll look into a fix
```
for k, v in pretrained_dict.items():
    print(type(k), type(v))

<class 'str'> <class 'int'>
<class 'str'> <class 'int'>
<class 'str'> <class 'str'>
<class 'str'> <class 'collections.OrderedDict'>
<class 'str'> <class 'str'>
<class 'str'> pytorch_lightning.utilities.parsing.AttributeDict
```

Edit 2:
I'll submit a PR after refactor week!

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
--- a/pytorch_lightning/trainer/training_io.py
+++ b/pytorch_lightning/trainer/training_io.py
@@ -370,10 +370,12 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
             if hasattr(model, '_hparams_name'):
                 checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_NAME] = model._hparams_name
             # add arguments to the checkpoint
-            checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY] = model.hparams
             if OMEGACONF_AVAILABLE:
+                checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY] = model.hparams
                 if isinstance(model.hparams, Container):
                     checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_TYPE] = type(model.hparams)
+            else:
+                checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY] = dict(model.hparams)
 
         # give the model a chance to add a few things
         model.on_save_checkpoint(checkpoint)


Test patch: