Improve collision check on hparams between LightningModule and DataModule (#9496)

* fix hyperparameter logging between LightningModule and DataModule

Author: ananthsub

CHANGELOG.md

@@ -204,6 +204,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Executing the `optimizer_closure` is now required when overriding the `optimizer_step` hook ([#9360](https://github.com/PyTorchLightning/pytorch-lightning/pull/9360))
 
 
+- Changed logging of `LightningModule` and `LightningDataModule` hyperparameters to raise an exception only if there are colliding keys with different values ([#9496](https://github.com/PyTorchLightning/pytorch-lightning/pull/9496))
+
+
 ### Deprecated
 
 - Deprecated `LightningModule.summarize()` in favor of `pytorch_lightning.utilities.model_summary.summarize()`

pytorch_lightning/trainer/trainer.py

@@ -1036,7 +1036,7 @@ def _pre_dispatch(self):
         self.accelerator.pre_dispatch(self)
         self._log_hyperparams()
 
-    def _log_hyperparams(self):
+    def _log_hyperparams(self) -> None:
         # log hyper-parameters
         hparams_initial = None
 
@@ -1047,12 +1047,20 @@ def _log_hyperparams(self):
             if self.lightning_module._log_hyperparams and datamodule_log_hyperparams:
                 datamodule_hparams = self.datamodule.hparams_initial
                 lightning_hparams = self.lightning_module.hparams_initial
-
-                colliding_keys = lightning_hparams.keys() & datamodule_hparams.keys()
-                if colliding_keys:
+                inconsistent_keys = []
+                for key in lightning_hparams.keys() & datamodule_hparams.keys():
+                    lm_val, dm_val = lightning_hparams[key], datamodule_hparams[key]
+                    if type(lm_val) != type(dm_val):
+                        inconsistent_keys.append(key)
+                    elif isinstance(lm_val, torch.Tensor) and id(lm_val) != id(dm_val):
+                        inconsistent_keys.append(key)
+                    elif lm_val != dm_val:
+                        inconsistent_keys.append(key)
+                if inconsistent_keys:
                     raise MisconfigurationException(
-                        f"Error while merging hparams: the keys {colliding_keys} are present "
-                        "in both the LightningModule's and LightningDataModule's hparams."
+                        f"Error while merging hparams: the keys {inconsistent_keys} are present "
+                        "in both the LightningModule's and LightningDataModule's hparams "
+                        "but have different values."
                     )
                 hparams_initial = {**lightning_hparams, **datamodule_hparams}
             elif self.lightning_module._log_hyperparams:

tests/loggers/test_base.py

@@ -13,17 +13,20 @@
 # limitations under the License.
 import pickle
 from argparse import Namespace
-from typing import Optional
+from copy import deepcopy
+from typing import Any, Dict, Optional
 from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
+import torch
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import LightningLoggerBase, LoggerCollection, TensorBoardLogger
 from pytorch_lightning.loggers.base import DummyExperiment, DummyLogger
 from pytorch_lightning.utilities import rank_zero_only
-from tests.helpers import BoringModel
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from tests.helpers.boring_model import BoringDataModule, BoringModel
 
 
 def test_logger_collection():
@@ -288,3 +291,77 @@ def __init__(self, param_one, param_two):
         log_hyperparams_mock.assert_called()
     else:
         log_hyperparams_mock.assert_not_called()
+
+
+@patch("pytorch_lightning.loggers.tensorboard.TensorBoardLogger.log_hyperparams")
+def test_log_hyperparams_key_collision(log_hyperparams_mock, tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self, hparams: Dict[str, Any]) -> None:
+            super().__init__()
+            self.save_hyperparameters(hparams)
+
+    class TestDataModule(BoringDataModule):
+        def __init__(self, hparams: Dict[str, Any]) -> None:
+            super().__init__()
+            self.save_hyperparameters(hparams)
+
+    class _Test:
+        ...
+
+    same_params = {1: 1, "2": 2, "three": 3.0, "test": _Test(), "4": torch.tensor(4)}
+    model = TestModel(same_params)
+    dm = TestDataModule(same_params)
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        limit_train_batches=0.1,
+        limit_val_batches=0.1,
+        num_sanity_val_steps=0,
+        checkpoint_callback=False,
+        progress_bar_refresh_rate=0,
+        weights_summary=None,
+    )
+    # there should be no exceptions raised for the same key/value pair in the hparams of both
+    # the lightning module and data module
+    trainer.fit(model)
+
+    obj_params = deepcopy(same_params)
+    obj_params["test"] = _Test()
+    model = TestModel(same_params)
+    dm = TestDataModule(obj_params)
+    trainer.fit(model)
+
+    diff_params = deepcopy(same_params)
+    diff_params.update({1: 0, "test": _Test()})
+    model = TestModel(same_params)
+    dm = TestDataModule(diff_params)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        limit_train_batches=0.1,
+        limit_val_batches=0.1,
+        num_sanity_val_steps=0,
+        checkpoint_callback=False,
+        progress_bar_refresh_rate=0,
+        weights_summary=None,
+    )
+    with pytest.raises(MisconfigurationException, match="Error while merging hparams"):
+        trainer.fit(model, dm)
+
+    tensor_params = deepcopy(same_params)
+    tensor_params.update({"4": torch.tensor(3)})
+    model = TestModel(same_params)
+    dm = TestDataModule(tensor_params)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        limit_train_batches=0.1,
+        limit_val_batches=0.1,
+        num_sanity_val_steps=0,
+        checkpoint_callback=False,
+        progress_bar_refresh_rate=0,
+        weights_summary=None,
+    )
+    with pytest.raises(MisconfigurationException, match="Error while merging hparams"):
+        trainer.fit(model, dm)