Fix LearningRateMonitor logging with multiple param groups optimizer with no scheduler (#10044)

Author: rohitgr7

CHANGELOG.md

@@ -596,6 +596,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `train_dataloader` getting loaded twice when resuming from a checkpoint during `Trainer.fit()` ([#9671](https://github.com/PyTorchLightning/pytorch-lightning/pull/9671))
 
 
+- Fixed `LearningRateMonitor` logging with multiple param groups optimizer with no scheduler ([#10044](https://github.com/PyTorchLightning/pytorch-lightning/pull/10044))
+
+
 
 - Fixed undesired side effects being caused by `Trainer` patching dataloader methods on the `LightningModule` ([#9764](https://github.com/PyTorchLightning/pytorch-lightning/pull/9764))
 

pytorch_lightning/callbacks/lr_monitor.py

@@ -19,6 +19,7 @@
 Monitor and logs learning rate for lr schedulers during training.
 
 """
+import itertools
 from collections import defaultdict
 from typing import Any, DefaultDict, Dict, List, Optional, Set, Tuple, Type
 
@@ -123,7 +124,7 @@ def _check_no_key(key: str) -> bool:
                 )
 
         # Find names for schedulers
-        names: List[str] = []
+        names: List[List[str]] = []
         (
             sched_hparam_keys,
             optimizers_with_scheduler,
@@ -140,8 +141,9 @@ def _check_no_key(key: str) -> bool:
         names.extend(optimizer_hparam_keys)
 
         # Initialize for storing values
-        self.lrs = {name: [] for name in names}
-        self.last_momentum_values = {name + "-momentum": None for name in names}
+        names_flatten = list(itertools.chain.from_iterable(names))
+        self.lrs = {name: [] for name in names_flatten}
+        self.last_momentum_values = {name + "-momentum": None for name in names_flatten}
 
     def on_train_batch_start(self, trainer: "pl.Trainer", *args: Any, **kwargs: Any) -> None:
         if not trainer.logger_connector.should_update_logs:
@@ -172,7 +174,7 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
         ) = self._find_names_from_schedulers(trainer.lr_schedulers, add_lr_sch_names=False)
         self._remap_keys(scheduler_hparam_keys)
 
-        for name, scheduler in zip(self.lr_sch_names, trainer.lr_schedulers):
+        for name, scheduler in zip(scheduler_hparam_keys, trainer.lr_schedulers):
             if interval in [scheduler["interval"], "any"]:
                 opt = scheduler["scheduler"].optimizer
                 current_stat = self._get_lr_momentum_stat(opt, name)
@@ -186,23 +188,22 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
         )
         self._remap_keys(optimizer_hparam_keys)
 
-        for opt, name in zip(optimizers_without_scheduler, optimizer_hparam_keys):
-            current_stat = self._get_lr_momentum_stat(opt, name)
+        for opt, names in zip(optimizers_without_scheduler, optimizer_hparam_keys):
+            current_stat = self._get_lr_momentum_stat(opt, names)
             latest_stat.update(current_stat)
 
         return latest_stat
 
-    def _get_lr_momentum_stat(self, optimizer: Optimizer, name: str) -> Dict[str, float]:
+    def _get_lr_momentum_stat(self, optimizer: Optimizer, names: List[str]) -> Dict[str, float]:
         lr_momentum_stat = {}
         param_groups = optimizer.param_groups
         use_betas = "betas" in optimizer.defaults
 
-        for i, pg in enumerate(param_groups):
-            name_and_suffix = self._add_suffix(name, param_groups, i)
-            lr = self._extract_lr(pg, name_and_suffix)
+        for pg, name in zip(param_groups, names):
+            lr = self._extract_lr(pg, name)
             lr_momentum_stat.update(lr)
             momentum = self._extract_momentum(
-                param_group=pg, name=name_and_suffix.replace(name, f"{name}-momentum"), use_betas=use_betas
+                param_group=pg, name=name.replace(name, f"{name}-momentum"), use_betas=use_betas
             )
             lr_momentum_stat.update(momentum)
 
@@ -213,14 +214,15 @@ def _extract_lr(self, param_group: Dict[str, Any], name: str) -> Dict[str, Any]:
         self.lrs[name].append(lr)
         return {name: lr}
 
-    def _remap_keys(self, names: List[str], token: str = "/pg1") -> None:
+    def _remap_keys(self, names: List[List[str]], token: str = "/pg1") -> None:
         """This function is used the remap the keys if param groups for a given optimizer increased."""
-        for new_name in names:
-            old_name = new_name.replace(token, "")
-            if token in new_name and old_name in self.lrs:
-                self.lrs[new_name] = self.lrs.pop(old_name)
-            elif new_name not in self.lrs:
-                self.lrs[new_name] = []
+        for group_new_names in names:
+            for new_name in group_new_names:
+                old_name = new_name.replace(token, "")
+                if token in new_name and old_name in self.lrs:
+                    self.lrs[new_name] = self.lrs.pop(old_name)
+                elif new_name not in self.lrs:
+                    self.lrs[new_name] = []
 
     def _extract_momentum(self, param_group: Dict[str, List], name: str, use_betas: bool) -> Dict[str, float]:
         if not self.log_momentum:
@@ -258,7 +260,7 @@ def _duplicate_param_group_names(self, param_groups: List[Dict]) -> Set[str]:
 
     def _find_names_from_schedulers(
         self, lr_schedulers: List, add_lr_sch_names: bool = True
-    ) -> Tuple[List[str], List[Optimizer], DefaultDict[Type[Optimizer], int]]:
+    ) -> Tuple[List[List[str]], List[Optimizer], DefaultDict[Type[Optimizer], int]]:
         # Create unique names in the case we have multiple of the same learning
         # rate scheduler + multiple parameter groups
         names = []
@@ -271,10 +273,11 @@ def _find_names_from_schedulers(
             else:
                 name = "lr-" + sch.optimizer.__class__.__name__
 
-            updated_name = self._check_duplicates_and_update_name(
+            updated_names = self._check_duplicates_and_update_name(
                 sch.optimizer, name, seen_optimizers, seen_optimizer_types, scheduler, add_lr_sch_names
             )
-            names.extend(updated_name)
+            names.append(updated_names)
+
         return names, seen_optimizers, seen_optimizer_types
 
     def _find_names_from_optimizers(
@@ -283,7 +286,7 @@ def _find_names_from_optimizers(
         seen_optimizers: List[Optimizer],
         seen_optimizer_types: DefaultDict[Type[Optimizer], int],
         add_lr_sch_names: bool = True,
-    ) -> Tuple[List[str], List[Optimizer]]:
+    ) -> Tuple[List[List[str]], List[Optimizer]]:
         names = []
         optimizers_without_scheduler = []
 
@@ -294,11 +297,12 @@ def _find_names_from_optimizers(
                 continue
 
             name = "lr-" + optimizer.__class__.__name__
-            updated_name = self._check_duplicates_and_update_name(
+            updated_names = self._check_duplicates_and_update_name(
                 optimizer, name, seen_optimizers, seen_optimizer_types, None, add_lr_sch_names
             )
-            names.extend(updated_name)
+            names.append(updated_names)
             optimizers_without_scheduler.append(optimizer)
+
         return names, optimizers_without_scheduler
 
     def _check_duplicates_and_update_name(

tests/callbacks/test_lr_monitor.py

@@ -510,3 +510,48 @@ def finetune_function(self, pl_module, epoch: int, optimizer, opt_idx: int):
 
     expected = [0.1, 0.05]
     assert lr_monitor.lrs["lr-Adam-1/pg3"] == expected
+
+
+def test_lr_monitor_multiple_param_groups_no_scheduler(tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self, lr, momentum):
+            super().__init__()
+            self.save_hyperparameters()
+            self.linear_a = torch.nn.Linear(32, 16)
+            self.linear_b = torch.nn.Linear(16, 2)
+
+        def forward(self, x):
+            x = self.linear_a(x)
+            x = self.linear_b(x)
+            return x
+
+        def configure_optimizers(self):
+            param_groups = [
+                {"params": list(self.linear_a.parameters())},
+                {"params": list(self.linear_b.parameters())},
+            ]
+            optimizer = torch.optim.Adam(param_groups, lr=self.hparams.lr, betas=self.hparams.momentum)
+            return optimizer
+
+    lr_monitor = LearningRateMonitor(log_momentum=True)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=2,
+        limit_val_batches=2,
+        limit_train_batches=2,
+        callbacks=[lr_monitor],
+        enable_progress_bar=False,
+        enable_model_summary=False,
+    )
+
+    lr = 1e-2
+    momentum = 0.7
+    model = TestModel(lr=lr, momentum=(momentum, 0.999))
+    trainer.fit(model)
+
+    assert len(lr_monitor.lrs) == len(trainer.optimizers[0].param_groups)
+    assert list(lr_monitor.lrs.keys()) == ["lr-Adam/pg1", "lr-Adam/pg2"]
+    assert lr_monitor.lr_sch_names == ["lr-Adam"]
+    assert list(lr_monitor.last_momentum_values.keys()) == ["lr-Adam/pg1-momentum", "lr-Adam/pg2-momentum"]
+    assert all(val == momentum for val in lr_monitor.last_momentum_values.values())
+    assert all(all(val == lr for val in lr_monitor.lrs[lr_key]) for lr_key in lr_monitor.lrs)