Repository: Lightning-AI/pytorch-lightning

Issue:
self.hparam silently removes params that are not serializable
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Following the approach found under [hyperparameters in the docs](https://pytorch-lightning.readthedocs.io/en/latest/hyperparameters.html#lightningmodule-hyperparameters), step 3, I passed a dict with parameters to my `pl.LightningModule`.
In the `__init__` printing `self.hparams` shows all contents I passed.
However, in the function `def configure_optimizers(self)`, some hparams are gone.

This might be related to that not all param values are YAML serializable, and therefore automatically removed? Because the 2 removed params are `"criterion": torch.nn.BCELoss()` and `"optimizer": partial(optim.Adam, lr=0.001)`.

### To Reproduce

Steps to reproduce the behavior:

1. Run the following script:
```python
from functools import partial
import torch
import torch.optim as optim
from torch.utils.data import Dataset
import pytorch_lightning as pl
from pytorch_lightning import Trainer

# partial to give all params, except the data
hparams = {
    "criterion": torch.nn.BCELoss(),  # F.cross_entropy(),  # loss function
    "optimizer": partial(optim.Adam, lr=0.001),  # (lr=0.001),
    # "learning_rate": 0.001,
    "filters": 64,
    "layers": 2
}

class EmptyDataset(Dataset):
    def __init__(self, transform=None):
        pass
    
    def __len__(self):
        return 32
    
    def __getitem__(self, idx):
        return {"input": np.array([1]), "output": "nothing"}

class LitLake(pl.LightningModule):
    def __init__(self, hparams: dict, transforms: dict = None):
        super().__init__()
        
        self.hparams = hparams
        print("self.hparams\n", self.hparams)
        
    def forward(self, x):
        pass
    
    def training_step(self, batch, batch_idx):
        """
        Lightning calls this inside the training loop with the data from the training dataloader
        passed in as `batch`.
        """
        # forward pass
        x, y = batch
        y_hat = self(x)
        loss = self.hparams["criterion"](y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}
        
    def configure_optimizers(self):
        print("self.hparams\n", self.hparams)
        optimizer = self.hparams["optimizer"](self.parameters())
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
        return [optimizer], [scheduler]
    
    def train_dataloader(self):
         return DataLoader(EmptyDataset(), batch_size=4, num_workers=1)

model = LitLake(hparams=hparams)
# most basic trainer, uses good defaults
trainer = Trainer()  # gpus=1, num_nodes=1
trainer.fit(model)  # KeyError: 'optimizer'
```
2. See error

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details><summary>Script output (CLICK ME)</summary>
<p>

```python
self.hparams
 "criterion": BCELoss()
"filters":   64
"layers":    2
"optimizer": functools.partial(<class 'torch.optim.adam.Adam'>, lr=0.001)
GPU available: True, used: False
TPU available: False, using: 0 TPU cores
self.hparams
 "filters": 64
"layers":  2
Traceback (most recent call last):
  File "lightning_hparams_bug.py", line 61, in <module>
    trainer.fit(model)  # KeyError: 'optimizer'
  File "/home/*user*/anaconda3/envs/onseilake/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 965, in fit
    self.optimizers, self.lr_schedulers, self.optimizer_frequencies = self.init_optimizers(model)
  File "/home/*user*/anaconda3/envs/onseilake/lib/python3.7/site-packages/pytorch_lightning/trainer/optimizers.py", line 18, in init_optimizers
    optim_conf = model.configure_optimizers()
  File "lightning_hparams_bug.py", line 51, in configure_optimizers
    optimizer = self.hparams["optimizer"](self.parameters())
KeyError: 'optimizer'
```

</p>
</details>

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
Either:
1. `self.hparams` keeping the non-serializable parameters (will give problems with loading?)
2. Throw an error explaining why those param values are not acceptable and how to approach it, instead of silently removing them.

### Environment

 - PyTorch Version (e.g., 1.0): 1.5.0
 - OS (e.g., Linux): Ubuntu 18.04
 - How you installed PyTorch (`conda`, `pip`, source): `conda`
 - Build command you used (if compiling from source): -
 - Python version: 3.7.6
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: 1x GeForce GTX 1080 Ti
 - Any other relevant information: -

### Additional context

<!-- Add any other context about the problem here. -->
Bug reproduced by @Borda 


Discussion:
Happens to me too. I'm also waiting for this fix, since i really want to use the new feature to put anything in hparams. 

@dscarmo you can take it over and send a PR :] or I ll check it tomorrow...
Need to add warning about this.

Resolving patch:
diff --git a/pytorch_lightning/utilities/__init__.py b/pytorch_lightning/utilities/__init__.py
--- a/pytorch_lightning/utilities/__init__.py
+++ b/pytorch_lightning/utilities/__init__.py
@@ -5,7 +5,7 @@
 
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 from pytorch_lightning.utilities.distributed import rank_zero_only, rank_zero_warn, rank_zero_info
-from pytorch_lightning.utilities.parsing import AttributeDict, flatten_dict
+from pytorch_lightning.utilities.parsing import AttributeDict, flatten_dict, is_picklable
 
 try:
     from apex import amp
diff --git a/pytorch_lightning/utilities/parsing.py b/pytorch_lightning/utilities/parsing.py
--- a/pytorch_lightning/utilities/parsing.py
+++ b/pytorch_lightning/utilities/parsing.py
@@ -1,7 +1,10 @@
 import inspect
+import pickle
 from argparse import Namespace
 from typing import Dict
 
+from pytorch_lightning.utilities import rank_zero_warn
+
 
 def str_to_bool(val):
     """Convert a string representation of truth to true (1) or false (0).
@@ -25,26 +28,28 @@ def str_to_bool(val):
         raise ValueError(f'invalid truth value {val}')
 
 
+def is_picklable(obj: object) -> bool:
+    """Tests if an object can be pickled"""
+
+    try:
+        pickle.dumps(obj)
+        return True
+    except pickle.PicklingError:
+        return False
+
+
 def clean_namespace(hparams):
-    """Removes all functions from hparams so we can pickle."""
+    """Removes all unpicklable entries from hparams"""
 
+    hparams_dict = hparams
     if isinstance(hparams, Namespace):
-        del_attrs = []
-        for k in hparams.__dict__:
-            if callable(getattr(hparams, k)):
-                del_attrs.append(k)
-
-        for k in del_attrs:
-            delattr(hparams, k)
-
-    elif isinstance(hparams, dict):
-        del_attrs = []
-        for k, v in hparams.items():
-            if callable(v):
-                del_attrs.append(k)
-
-        for k in del_attrs:
-            del hparams[k]
+        hparams_dict = hparams.__dict__
+
+    del_attrs = [k for k, v in hparams_dict.items() if not is_picklable(v)]
+
+    for k in del_attrs:
+        rank_zero_warn(f"attribute '{k}' removed from hparams because it cannot be pickled", UserWarning)
+        del hparams_dict[k]
 
 
 def get_init_args(frame) -> dict:


Test patch:
diff --git a/tests/models/test_hparams.py b/tests/models/test_hparams.py
--- a/tests/models/test_hparams.py
+++ b/tests/models/test_hparams.py
@@ -11,7 +11,7 @@
 
 from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.core.saving import save_hparams_to_yaml, load_hparams_from_yaml
-from pytorch_lightning.utilities import AttributeDict
+from pytorch_lightning.utilities import AttributeDict, is_picklable
 from tests.base import EvalModelTemplate, TrialMNIST
 
 
@@ -282,7 +282,7 @@ def test_collect_init_arguments(tmpdir, cls):
     assert model.hparams.batch_size == 179
 
     if isinstance(model, AggSubClassEvalModel):
-        assert isinstance(model.hparams.my_loss, torch.nn.CrossEntropyLoss)
+        assert isinstance(model.hparams.my_loss, torch.nn.CosineEmbeddingLoss)
 
     if isinstance(model, DictConfSubClassEvalModel):
         assert isinstance(model.hparams.dict_conf, Container)
@@ -413,6 +413,23 @@ def test_hparams_pickle(tmpdir):
     assert ad == pickle.loads(pkl)
 
 
+class UnpickleableArgsEvalModel(EvalModelTemplate):
+    """ A model that has an attribute that cannot be pickled. """
+
+    def __init__(self, foo='bar', pickle_me=(lambda x: x + 1), **kwargs):
+        super().__init__(**kwargs)
+        assert not is_picklable(pickle_me)
+        self.save_hyperparameters()
+
+
+def test_hparams_pickle_warning(tmpdir):
+    model = UnpickleableArgsEvalModel()
+    trainer = Trainer(default_root_dir=tmpdir, max_steps=1)
+    with pytest.warns(UserWarning, match="attribute 'pickle_me' removed from hparams because it cannot be pickled"):
+        trainer.fit(model)
+    assert 'pickle_me' not in model.hparams
+
+
 def test_hparams_save_yaml(tmpdir):
     hparams = dict(batch_size=32, learning_rate=0.001, data_root='./any/path/here',
                    nasted=dict(any_num=123, anystr='abcd'))