Repository: Lightning-AI/pytorch-lightning
Pull request: [bugfix] Resolve LearningRateMonitor + BackboneFinetuning
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->
Fixes #7830

## Before submitting
- [ ] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes? (if necessary)
- [ ] Did you write any new necessary tests? (not for typos and docs)
- [ ] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Error with `BackboneFinetuning` and `LearningRateMonitor` with `logging_interval="epoch"`
## 🐛 Bug

When using the backbone finetuning callback + learning rate moniter with logging interval = epoch an error is thrown after the backbone is unfrozen.

## Please reproduce using the BoringModel

```
import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks.finetuning import BackboneFinetuning
from pytorch_lightning.callbacks.lr_monitor import LearningRateMonitor


class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):

    def __init__(self):
        super().__init__()
        self.backbone = torch.nn.Sequential(torch.nn.Linear(32, 32), torch.nn.ReLU(True))
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(self.backbone(x))

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        opt = torch.optim.Adam(self.layer.parameters(), lr=0.1)
        return [opt], [torch.optim.lr_scheduler.StepLR(opt, step_size=1)]


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=6,
        weights_summary=None,
        callbacks=[BackboneFinetuning(unfreeze_backbone_at_epoch=4), LearningRateMonitor(logging_interval="epoch")]
    )
    trainer.fit(model, train_dataloader=train_data, val_dataloaders=val_data)
    trainer.test(model, test_dataloaders=test_data)


if __name__ == '__main__':
    run()
```

Error:
```
lr_monitor.py", line 160, in _extract_lr
    self.lrs[name].append(lr)
KeyError: 'lr-Adam/pg1'
Exception ignored in: <function tqdm.__del__ at 0x7ff5e83d7d30>
Traceback (most recent call last):
  File "/Users/ethan/miniconda3/lib/python3.8/site-packages/tqdm/std.py", line 1090, in __del__
  File "/Users/ethan/miniconda3/lib/python3.8/site-packages/tqdm/std.py", line 1303, in close
  File "/Users/ethan/miniconda3/lib/python3.8/site-packages/tqdm/std.py", line 1481, in display
  File "/Users/ethan/miniconda3/lib/python3.8/site-packages/tqdm/std.py", line 1093, in __repr__
  File "/Users/ethan/miniconda3/lib/python3.8/site-packages/tqdm/std.py", line 1443, in format_dict
TypeError: cannot unpack non-iterable NoneType object
```

### Expected behavior

No error

### Environment

PL master branch

Hints from issue discussion:

Commit message: update changelog

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f6d4692ec076..5fc69fb34a46e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -191,6 +191,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a bug where checking `trainer.precision` changed to `'mixed'` when specifying 16 in trainer ([#7825](https://github.com/PyTorchLightning/pytorch-lightning/pull/7825))
 
 
+- Fixed `LearningRateMonitor` keys not properly setup when running with `BackboneFinetuning` Callback ([#7835](https://github.com/PyTorchLightning/pytorch-lightning/pull/7835))
+
+
 ## [1.3.2] - 2021-05-18
 
 ### Changed
diff --git a/pytorch_lightning/callbacks/lr_monitor.py b/pytorch_lightning/callbacks/lr_monitor.py
index 410f8b319c239..5d2d5eae046aa 100644
--- a/pytorch_lightning/callbacks/lr_monitor.py
+++ b/pytorch_lightning/callbacks/lr_monitor.py
@@ -19,7 +19,7 @@
 Monitor and logs learning rate for lr schedulers during training.
 
 """
-
+from collections import defaultdict
 from typing import Dict, List, Optional
 
 from pytorch_lightning.callbacks.base import Callback
@@ -53,7 +53,7 @@ class LearningRateMonitor(Callback):
     In case of multiple optimizers of same type, they will be named ``Adam``,
     ``Adam-1`` etc. If a optimizer has multiple parameter groups they will
     be named ``Adam/pg1``, ``Adam/pg2`` etc. To control naming, pass in a
-    ``name`` keyword in the construction of the learning rate schdulers
+    ``name`` keyword in the construction of the learning rate schedulers
 
     Example::
 
@@ -146,7 +146,7 @@ def _extract_stats(self, trainer, interval: str) -> Dict[str, float]:
 
                 for i, pg in enumerate(param_groups):
                     suffix = f'/pg{i + 1}' if len(param_groups) > 1 else ''
-                    lr = self._extract_lr(param_group=pg, name=f'{name}{suffix}')
+                    lr = self._extract_lr(trainer, param_group=pg, name=f'{name}{suffix}')
                     latest_stat.update(lr)
                     momentum = self._extract_momentum(
                         param_group=pg, name=f'{name}-momentum{suffix}', use_betas=use_betas
@@ -155,11 +155,26 @@ def _extract_stats(self, trainer, interval: str) -> Dict[str, float]:
 
         return latest_stat
 
-    def _extract_lr(self, param_group, name: str) -> Dict[str, float]:
+    def _extract_lr(self, trainer, param_group, name: str) -> Dict[str, float]:
         lr = param_group.get('lr')
-        self.lrs[name].append(lr)
+        try:
+            self.lrs[name].append(lr)
+        except KeyError:
+            names = self._find_names(trainer.lr_schedulers, add_lr_sch_names=False)
+            self._remap_keys(names)
+            self.lrs[name].append(lr)
         return {name: lr}
 
+    def _remap_keys(self, names: List[str]) -> None:
+        token = '/pg1'
+        for new_name in names:
+            if token in new_name:
+                old_n = new_name.replace(token, '')
+                self.lrs[new_name] = self.lrs[old_n]
+                del self.lrs[old_n]
+            else:
+                self.lrs[new_name] = []
+
     def _extract_momentum(self, param_group, name: str, use_betas: bool) -> Dict[str, float]:
         if not self.log_momentum:
             return {}
@@ -168,35 +183,45 @@ def _extract_momentum(self, param_group, name: str, use_betas: bool) -> Dict[str
         self.last_momentum_values[name] = momentum
         return {name: momentum}
 
-    def _find_names(self, lr_schedulers) -> List[str]:
-        # Create uniqe names in the case we have multiple of the same learning
-        # rate schduler + multiple parameter groups
+    def _add_prefix(self, name, optimizer_cls, seen_optimizer_types) -> str:
+        count = seen_optimizer_types[optimizer_cls]
+        return name + f'-{count}' if count else name
+
+    def _find_names(self, lr_schedulers, add_lr_sch_names: bool = True) -> List[str]:
+        # Create unique names in the case we have multiple of the same learning
+        # rate scheduler + multiple parameter groups
         names = []
+        seen_optimizers = []
+        seen_optimizer_types = defaultdict(int)
         for scheduler in lr_schedulers:
             sch = scheduler['scheduler']
             if scheduler['name'] is not None:
                 name = scheduler['name']
             else:
-                opt_name = 'lr-' + sch.optimizer.__class__.__name__
-                i, name = 1, opt_name
+                name = 'lr-' + sch.optimizer.__class__.__name__
 
-                # Multiple schduler of the same type
-                while True:
-                    if name not in names:
-                        break
-                    i, name = i + 1, f'{opt_name}-{i}'
+            seen_optimizers.append(sch.optimizer)
+            optimizer_cls = type(sch.optimizer)
+            if scheduler['name'] is None:
+                if optimizer_cls not in seen_optimizer_types:
+                    seen_optimizer_types[optimizer_cls] = 0
+                else:
+                    seen_optimizer_types[optimizer_cls] += 1
 
-            # Multiple param groups for the same schduler
+            # Multiple param groups for the same scheduler
             param_groups = sch.optimizer.param_groups
 
             if len(param_groups) != 1:
-                for i, pg in enumerate(param_groups):
-                    temp = f'{name}/pg{i + 1}'
+                for i in range(len(param_groups)):
+                    temp = self._add_prefix(name, optimizer_cls, seen_optimizer_types)
+                    temp = f'{temp}/pg{i + 1}'
                     names.append(temp)
             else:
+                name = self._add_prefix(name, optimizer_cls, seen_optimizer_types)
                 names.append(name)
 
-            self.lr_sch_names.append(name)
+            if add_lr_sch_names:
+                self.lr_sch_names.append(name)
 
         return names
 
diff --git a/tests/callbacks/test_lr_monitor.py b/tests/callbacks/test_lr_monitor.py
index bea6c45e95ced..ac3939fc362ce 100644
--- a/tests/callbacks/test_lr_monitor.py
+++ b/tests/callbacks/test_lr_monitor.py
@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
+import torch
 from torch import optim
 
 import tests.helpers.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.callbacks.base import Callback
+from pytorch_lightning.callbacks.finetuning import BackboneFinetuning
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
 from tests.helpers.datamodules import ClassifDataModule
@@ -278,3 +281,54 @@ def configure_optimizers(self):
     )
     trainer.fit(TestModel())
     assert lr_monitor.lr_sch_names == list(lr_monitor.lrs.keys()) == ['my_logging_name']
+
+
+def test_multiple_optimizers_basefinetuning(tmpdir):
+
+    class TestModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.backbone = torch.nn.Sequential(torch.nn.Linear(32, 32), torch.nn.ReLU(True))
+            self.layer = torch.nn.Linear(32, 2)
+
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            return super().training_step(batch, batch_idx)
+
+        def forward(self, x):
+            return self.layer(self.backbone(x))
+
+        def configure_optimizers(self):
+            opt = optim.Adam(self.layer.parameters(), lr=0.1)
+            opt_2 = optim.Adam(self.layer.parameters(), lr=0.1)
+            opt_3 = optim.Adam(self.layer.parameters(), lr=0.1)
+            return [opt, opt_2, opt_3
+                    ], [optim.lr_scheduler.StepLR(opt, step_size=1),
+                        optim.lr_scheduler.StepLR(opt_2, step_size=1)]
+
+    class Check(Callback):
+
+        def on_epoch_end(self, trainer, pl_module) -> None:
+            assert lr_monitor.lr_sch_names == ['lr-Adam', 'lr-Adam-1']
+            if trainer.current_epoch > 2:
+                assert list(lr_monitor.lrs.keys()) == ['lr-Adam/pg1', 'lr-Adam/pg2', 'lr-Adam-1/pg1', 'lr-Adam-1/pg2']
+            else:
+                assert list(lr_monitor.lrs.keys()) == ['lr-Adam', 'lr-Adam-1']
+
+    lr_monitor = LearningRateMonitor()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=4,
+        limit_val_batches=0,
+        limit_train_batches=2,
+        callbacks=[lr_monitor, BackboneFinetuning(unfreeze_backbone_at_epoch=2),
+                   Check()],
+        progress_bar_refresh_rate=0,
+        weights_summary=None,
+    )
+    model = TestModel()
+    model.training_epoch_end = None
+    trainer.fit(model)
+
+    # 3 epoch difference
+    assert len(lr_monitor.lrs["lr-Adam/pg1"]) == len(lr_monitor.lrs["lr-Adam/pg2"]) + 3


Code review comments:
diff_hunk: @@ -278,3 +281,54 @@ def configure_optimizers(self):
     )
     trainer.fit(TestModel())
     assert lr_monitor.lr_sch_names == list(lr_monitor.lrs.keys()) == ['my_logging_name']
+
+
+def test_multiple_optimizers_basefinetuning(tmpdir):
+
+    class TestModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.backbone = torch.nn.Sequential(torch.nn.Linear(32, 32), torch.nn.ReLU(True))
+            self.layer = torch.nn.Linear(32, 2)
+
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            return super().training_step(batch, batch_idx)
+
+        def forward(self, x):
+            return self.layer(self.backbone(x))
+
+        def configure_optimizers(self):
+            opt = optim.Adam(self.layer.parameters(), lr=0.1)
+            opt_2 = optim.Adam(self.layer.parameters(), lr=0.1)
+            opt_3 = optim.Adam(self.layer.parameters(), lr=0.1)
+            return [opt, opt_2, opt_3
+                    ], [optim.lr_scheduler.StepLR(opt, step_size=1),
+                        optim.lr_scheduler.StepLR(opt_2, step_size=1)]
path: tests/callbacks/test_lr_monitor.py
text: @user1:
formatting issue? looks broken 😄 

@author:
Done !

diff_hunk: @@ -168,35 +183,45 @@ def _extract_momentum(self, param_group, name: str, use_betas: bool) -> Dict[str
         self.last_momentum_values[name] = momentum
         return {name: momentum}
 
-    def _find_names(self, lr_schedulers) -> List[str]:
-        # Create uniqe names in the case we have multiple of the same learning
-        # rate schduler + multiple parameter groups
+    def _add_prefix(self, name, optimizer_cls, seen_optimizer_types) -> str:
+        count = seen_optimizer_types[optimizer_cls]
+        return name + f'-{count}' if count else name
+
+    def _find_names(self, lr_schedulers, add_lr_sch_names: bool = True) -> List[str]:
+        # Create unique names in the case we have multiple of the same learning
+        # rate scheduler + multiple parameter groups
         names = []
+        seen_optimizers = []
+        seen_optimizer_types = defaultdict(int)
         for scheduler in lr_schedulers:
             sch = scheduler['scheduler']
             if scheduler['name'] is not None:
                 name = scheduler['name']
             else:
-                opt_name = 'lr-' + sch.optimizer.__class__.__name__
-                i, name = 1, opt_name
+                name = 'lr-' + sch.optimizer.__class__.__name__
 
-                # Multiple schduler of the same type
-                while True:
-                    if name not in names:
-                        break
-                    i, name = i + 1, f'{opt_name}-{i}'
+            seen_optimizers.append(sch.optimizer)
+            optimizer_cls = type(sch.optimizer)
+            if scheduler['name'] is None:
+                if optimizer_cls not in seen_optimizer_types:
path: pytorch_lightning/callbacks/lr_monitor.py
text: @user1:
As `seen_optimizer_types` is a `defaultdict` you should just be able to do `seen_optimizer_types[optimizer_cls] += 1` and not need the if. 

@author:
Done !

@user2:
Doesn't look done 🧐 

@author:
Added back !

diff_hunk: @@ -155,11 +155,26 @@ def _extract_stats(self, trainer, interval: str) -> Dict[str, float]:
 
         return latest_stat
 
-    def _extract_lr(self, param_group, name: str) -> Dict[str, float]:
+    def _extract_lr(self, trainer, param_group, name: str) -> Dict[str, float]:
         lr = param_group.get('lr')
-        self.lrs[name].append(lr)
+        try:
+            self.lrs[name].append(lr)
+        except KeyError:
path: pytorch_lightning/callbacks/lr_monitor.py
text: @user1:
it feels like `if name in self.lrs ... else ...` is better than try except but I could be wrong. maybe this is even slightly faster? ^^

@author:
Thanks @user1. Found a bug :)