Repository: Lightning-Universe/lightning-flash

Issue:
OneCycleLR scheduler does not work with freeze-unfreeze finetuning strategy
## 🐛 Bug

I wanted to create an image classifier by fine-tuning pre-trained model on my dataset. When OneCycleLR scheduler is used alongside the freeze-unfreeze, training throws an exception once the unfreeze epoch is reached.


### To Reproduce / Code Sample

I use flash's built-in `ImageClassifier` as follows:

```python
   epochs = 50
   model = ImageClassifier(
        backbone="efficientnet_b5",
        labels=datamodule.labels,
        metrics=[
            Accuracy(),
        ],
        optimizer="AdamW",
        lr_scheduler=(
            "onecyclelr",
            {
                "max_lr": 1e-3,
                "epochs": epochs,
                "steps_per_epoch": steps_per_epoch,
            },
            {"interval": "step"},
        ),
    )
```

```python
trainer = flash.Trainer(
    max_epochs=epochs,
    gpus=torch.cuda.device_count(),
)
trainer.finetune(model, datamodule=datamodule, strategy=("freeze_unfreeze", 5))
```

### Expected behaviour

After specified number of epochs, layers get unfrozen and training continues.

### Actual behaviour
Expection is thrown:

```python
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training
    self._results = trainer.run_stage()
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage
    return self._run_train()
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1319, in _run_train
    self.fit_loop.run()
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 234, in advance
    self.epoch_loop.run(data_fetcher)
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 199, in advance
    self.update_lr_schedulers("step", update_plateau_schedulers=False)
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 441, in update_lr_schedulers
    self._update_learning_rates(
  File "/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 505, in _update_learning_rates
    lr_scheduler["scheduler"].step()
  File "/<redacted>/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 154, in step
    values = self.get_lr()
  File "/<redacted>/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 1597, in get_lr
    computed_lr = self.anneal_func(group[phase['start_lr']], group[phase['end_lr']], pct)
KeyError: 'max_lr
```

It seems like the unfreezing strategy creates additional optimizer parameter groups, but when the unfreezing happens, some of the LR scheduler parameters are not copied / passed to the new param group properly in: `pytorch_lightning.callbacks.finetuning.BaseFinetuning.unfreeze_and_add_param_group`.

### Environment

 - OS (e.g., Linux): macOS
 - Python version: 3.8.12
 - PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7): 1.11.0 / 1.5.10 / 0.7.3
 - GPU models and configuration: 0 / 1 T4 (happens regardless of cuda)
 - Any other relevant information:

### Additional context

https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1651218144224359




Discussion:

Resolving patch:
diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py
--- a/flash/core/finetuning.py
+++ b/flash/core/finetuning.py
@@ -103,6 +103,19 @@ def freeze_before_training(self, pl_module: Union[Module, Iterable[Union[Module,
                     modules = [modules]
                 self.freeze(modules=modules, train_bn=self.train_bn)
 
+    def unfreeze_and_extend_param_group(
+        self,
+        modules: Union[Module, Iterable[Union[Module, Iterable]]],
+        optimizer: Optimizer,
+        train_bn: bool = True,
+    ) -> None:
+        self.make_trainable(modules)
+
+        params = self.filter_params(modules, train_bn=train_bn, requires_grad=True)
+        params = self.filter_on_optimizer(optimizer, params)
+        if params:
+            optimizer.param_groups[0]["params"].extend(params)
+
     def _freeze_unfreeze_function(
         self,
         pl_module: Union[Module, Iterable[Union[Module, Iterable]]],
@@ -117,7 +130,7 @@ def _freeze_unfreeze_function(
 
         modules = self._get_modules_to_freeze(pl_module=pl_module)
         if modules is not None:
-            self.unfreeze_and_add_param_group(
+            self.unfreeze_and_extend_param_group(
                 modules=modules,
                 optimizer=optimizer,
                 train_bn=self.train_bn,
@@ -140,7 +153,7 @@ def _unfreeze_milestones_function(
                 # unfreeze num_layers last layers
 
                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[-num_layers:]
-                self.unfreeze_and_add_param_group(
+                self.unfreeze_and_extend_param_group(
                     modules=backbone_modules,
                     optimizer=optimizer,
                     train_bn=self.train_bn,
@@ -148,7 +161,7 @@ def _unfreeze_milestones_function(
             elif epoch == unfreeze_milestones[1]:
                 # unfreeze remaining layers
                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[:-num_layers]
-                self.unfreeze_and_add_param_group(
+                self.unfreeze_and_extend_param_group(
                     modules=backbone_modules,
                     optimizer=optimizer,
                     train_bn=self.train_bn,


Test patch:
diff --git a/tests/core/test_finetuning.py b/tests/core/test_finetuning.py
--- a/tests/core/test_finetuning.py
+++ b/tests/core/test_finetuning.py
@@ -155,20 +155,45 @@ def test_finetuning_with_none_return_type(strategy):
 
 @pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.")
 @pytest.mark.parametrize(
-    ("strategy", "checker_class", "checker_class_data"),
+    ("strategy", "lr_scheduler", "checker_class", "checker_class_data"),
     [
-        ("no_freeze", None, {}),
-        ("freeze", FreezeStrategyChecking, {}),
-        (("freeze_unfreeze", 2), FreezeUnfreezeStrategyChecking, {"check_epoch": 2}),
+        ("no_freeze", None, None, {}),
+        ("freeze", None, FreezeStrategyChecking, {}),
+        (("freeze_unfreeze", 2), None, FreezeUnfreezeStrategyChecking, {"check_epoch": 2}),
         (
             ("unfreeze_milestones", ((1, 3), 1)),
+            None,
+            UnfreezeMilestonesStrategyChecking,
+            {"check_epochs": [1, 3], "num_layers": 1},
+        ),
+        (
+            "no_freeze",
+            ("onecyclelr", {"max_lr": 1e-3, "epochs": 50, "steps_per_epoch": 10}, {"interval": "step"}),
+            None,
+            {},
+        ),
+        (
+            "freeze",
+            ("onecyclelr", {"max_lr": 1e-3, "epochs": 50, "steps_per_epoch": 10}, {"interval": "step"}),
+            FreezeStrategyChecking,
+            {},
+        ),
+        (
+            ("freeze_unfreeze", 2),
+            ("onecyclelr", {"max_lr": 1e-3, "epochs": 50, "steps_per_epoch": 10}, {"interval": "step"}),
+            FreezeUnfreezeStrategyChecking,
+            {"check_epoch": 2},
+        ),
+        (
+            ("unfreeze_milestones", ((1, 3), 1)),
+            ("onecyclelr", {"max_lr": 1e-3, "epochs": 50, "steps_per_epoch": 10}, {"interval": "step"}),
             UnfreezeMilestonesStrategyChecking,
             {"check_epochs": [1, 3], "num_layers": 1},
         ),
     ],
 )
-def test_finetuning(tmpdir, strategy, checker_class, checker_class_data):
-    task = TestTaskWithFinetuning(loss_fn=F.nll_loss)
+def test_finetuning(tmpdir, strategy, lr_scheduler, checker_class, checker_class_data):
+    task = TestTaskWithFinetuning(loss_fn=F.nll_loss, lr_scheduler=lr_scheduler, optimizer="sgd", learning_rate=0.1)
     callbacks = [] if checker_class is None else checker_class(dirpath=tmpdir, **checker_class_data)
     trainer = flash.Trainer(max_epochs=5, limit_train_batches=10, callbacks=callbacks)
     ds = DummyDataset()