We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flash/core/finetuning.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 import os
15 from functools import partial
16 from typing import Iterable, Optional, Tuple, Union
17 
18 from pytorch_lightning import LightningModule
19 from pytorch_lightning.callbacks import BaseFinetuning
20 from pytorch_lightning.utilities.exceptions import MisconfigurationException
21 from torch.nn import Module
22 from torch.optim import Optimizer
23 
24 from flash.core.registry import FlashRegistry
25 
26 if not os.environ.get("READTHEDOCS", False):
27     from pytorch_lightning.utilities.enums import LightningEnum
28 else:
29     # ReadTheDocs mocks the `LightningEnum` import to be a regular type, so we replace it with a plain Enum here.
30     from enum import Enum
31 
32     LightningEnum = Enum
33 
34 
35 class FinetuningStrategies(LightningEnum):
36     """The ``FinetuningStrategies`` enum contains the keys that are used internally by the ``FlashBaseFinetuning``
37     when choosing the strategy to perform."""
38 
39     NO_FREEZE = "no_freeze"
40     FREEZE = "freeze"
41     FREEZE_UNFREEZE = "freeze_unfreeze"
42     UNFREEZE_MILESTONES = "unfreeze_milestones"
43 
44     # TODO: Create a FlashEnum class???
45     def __hash__(self) -> int:
46         return hash(self.value)
47 
48 
49 class FlashBaseFinetuning(BaseFinetuning):
50     """FlashBaseFinetuning can be used to create a custom Flash Finetuning Callback."""
51 
52     def __init__(
53         self,
54         strategy_key: Union[str, FinetuningStrategies],
55         strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = None,
56         train_bn: bool = True,
57     ):
58         """
59         Args:
60             strategy_key: The finetuning strategy to be used. See :meth:`~flash.core.trainer.Trainer.finetune`
61                 for the available strategies.
62             strategy_metadata: Data that accompanies certain finetuning strategies like epoch number or number of
63                 layers.
64             train_bn: Whether to train Batch Norm layer
65         """
66         super().__init__()
67 
68         self.strategy: FinetuningStrategies = strategy_key
69         self.strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = strategy_metadata
70         self.train_bn: bool = train_bn
71 
72         if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE and not isinstance(self.strategy_metadata, int):
73             raise MisconfigurationException(
74                 "The `freeze_unfreeze` strategy requires an integer denoting the epoch number to unfreeze at. Example: "
75                 "`strategy=('freeze_unfreeze', 7)`"
76             )
77         if self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES and not (
78             isinstance(self.strategy_metadata, Tuple)
79             and isinstance(self.strategy_metadata[0], Tuple)
80             and isinstance(self.strategy_metadata[1], int)
81             and isinstance(self.strategy_metadata[0][0], int)
82             and isinstance(self.strategy_metadata[0][1], int)
83         ):
84             raise MisconfigurationException(
85                 "The `unfreeze_milestones` strategy requires the format Tuple[Tuple[int, int], int]. Example: "
86                 "`strategy=('unfreeze_milestones', ((5, 10), 15))`"
87             )
88 
89     def _get_modules_to_freeze(self, pl_module: LightningModule) -> Union[Module, Iterable[Union[Module, Iterable]]]:
90         modules_to_freeze = getattr(pl_module, "modules_to_freeze", None)
91         if modules_to_freeze is None:
92             raise AttributeError(
93                 "LightningModule missing instance method 'modules_to_freeze'."
94                 "Please, implement the method which returns NoneType or a Module or an Iterable of Modules."
95             )
96         return modules_to_freeze()
97 
98     def freeze_before_training(self, pl_module: Union[Module, Iterable[Union[Module, Iterable]]]) -> None:
99         if self.strategy != FinetuningStrategies.NO_FREEZE:
100             modules = self._get_modules_to_freeze(pl_module=pl_module)
101             if modules is not None:
102                 if isinstance(modules, Module):
103                     modules = [modules]
104                 self.freeze(modules=modules, train_bn=self.train_bn)
105 
106     def _freeze_unfreeze_function(
107         self,
108         pl_module: Union[Module, Iterable[Union[Module, Iterable]]],
109         epoch: int,
110         optimizer: Optimizer,
111         opt_idx: int,
112         strategy_metadata: int,
113     ):
114         unfreeze_epoch: int = strategy_metadata
115         if epoch != unfreeze_epoch:
116             return
117 
118         modules = self._get_modules_to_freeze(pl_module=pl_module)
119         if modules is not None:
120             self.unfreeze_and_add_param_group(
121                 modules=modules,
122                 optimizer=optimizer,
123                 train_bn=self.train_bn,
124             )
125 
126     def _unfreeze_milestones_function(
127         self,
128         pl_module: Union[Module, Iterable[Union[Module, Iterable]]],
129         epoch: int,
130         optimizer: Optimizer,
131         opt_idx: int,
132         strategy_metadata: Tuple[Tuple[int, int], int],
133     ):
134         unfreeze_milestones: Tuple[int, int] = strategy_metadata[0]
135         num_layers: int = strategy_metadata[1]
136 
137         modules = self._get_modules_to_freeze(pl_module=pl_module)
138         if modules is not None:
139             if epoch == unfreeze_milestones[0]:
140                 # unfreeze num_layers last layers
141 
142                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[-num_layers:]
143                 self.unfreeze_and_add_param_group(
144                     modules=backbone_modules,
145                     optimizer=optimizer,
146                     train_bn=self.train_bn,
147                 )
148             elif epoch == unfreeze_milestones[1]:
149                 # unfreeze remaining layers
150                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[:-num_layers]
151                 self.unfreeze_and_add_param_group(
152                     modules=backbone_modules,
153                     optimizer=optimizer,
154                     train_bn=self.train_bn,
155                 )
156 
157     def finetune_function(
158         self,
159         pl_module: Union[Module, Iterable[Union[Module, Iterable]]],
160         epoch: int,
161         optimizer: Optimizer,
162         opt_idx: int,
163     ):
164         if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE:
165             self._freeze_unfreeze_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)
166         elif self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES:
167             self._unfreeze_milestones_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)
168 
169 
170 _FINETUNING_STRATEGIES_REGISTRY = FlashRegistry("finetuning_strategies")
171 
172 for strategy in FinetuningStrategies:
173     _FINETUNING_STRATEGIES_REGISTRY(
174         name=strategy.value,
175         fn=partial(FlashBaseFinetuning, strategy_key=strategy),
176     )
177 
178 
179 class NoFreeze(FlashBaseFinetuning):
180     def __init__(self, train_bn: bool = True):
181         super().__init__(FinetuningStrategies.NO_FREEZE, train_bn)
182 
183 
184 class Freeze(FlashBaseFinetuning):
185     def __init__(self, train_bn: bool = True):
186         super().__init__(FinetuningStrategies.FREEZE, train_bn)
187 
188 
189 class FreezeUnfreeze(FlashBaseFinetuning):
190     def __init__(
191         self,
192         strategy_metadata: int,
193         train_bn: bool = True,
194     ):
195         super().__init__(FinetuningStrategies.FREEZE_UNFREEZE, strategy_metadata, train_bn)
196 
197 
198 class UnfreezeMilestones(FlashBaseFinetuning):
199     def __init__(
200         self,
201         strategy_metadata: Tuple[Tuple[int, int], int],
202         train_bn: bool = True,
203     ):
204         super().__init__(FinetuningStrategies.UNFREEZE_MILESTONES, strategy_metadata, train_bn)
205
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
                     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 @@
 
         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 @@
                 # 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 @@
             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,

{"golden_diff": "diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py\n--- a/flash/core/finetuning.py\n+++ b/flash/core/finetuning.py\n@@ -103,6 +103,19 @@\n                     modules = [modules]\n                 self.freeze(modules=modules, train_bn=self.train_bn)\n \n+    def unfreeze_and_extend_param_group(\n+        self,\n+        modules: Union[Module, Iterable[Union[Module, Iterable]]],\n+        optimizer: Optimizer,\n+        train_bn: bool = True,\n+    ) -> None:\n+        self.make_trainable(modules)\n+\n+        params = self.filter_params(modules, train_bn=train_bn, requires_grad=True)\n+        params = self.filter_on_optimizer(optimizer, params)\n+        if params:\n+            optimizer.param_groups[0][\"params\"].extend(params)\n+\n     def _freeze_unfreeze_function(\n         self,\n         pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n@@ -117,7 +130,7 @@\n \n         modules = self._get_modules_to_freeze(pl_module=pl_module)\n         if modules is not None:\n-            self.unfreeze_and_add_param_group(\n+            self.unfreeze_and_extend_param_group(\n                 modules=modules,\n                 optimizer=optimizer,\n                 train_bn=self.train_bn,\n@@ -140,7 +153,7 @@\n                 # unfreeze num_layers last layers\n \n                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[-num_layers:]\n-                self.unfreeze_and_add_param_group(\n+                self.unfreeze_and_extend_param_group(\n                     modules=backbone_modules,\n                     optimizer=optimizer,\n                     train_bn=self.train_bn,\n@@ -148,7 +161,7 @@\n             elif epoch == unfreeze_milestones[1]:\n                 # unfreeze remaining layers\n                 backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[:-num_layers]\n-                self.unfreeze_and_add_param_group(\n+                self.unfreeze_and_extend_param_group(\n                     modules=backbone_modules,\n                     optimizer=optimizer,\n                     train_bn=self.train_bn,\n", "issue": "OneCycleLR scheduler does not work with freeze-unfreeze finetuning strategy\n## \ud83d\udc1b Bug\r\n\r\nI 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.\r\n\r\n\r\n### To Reproduce / Code Sample\r\n\r\nI use flash's built-in `ImageClassifier` as follows:\r\n\r\n```python\r\n   epochs = 50\r\n   model = ImageClassifier(\r\n        backbone=\"efficientnet_b5\",\r\n        labels=datamodule.labels,\r\n        metrics=[\r\n            Accuracy(),\r\n        ],\r\n        optimizer=\"AdamW\",\r\n        lr_scheduler=(\r\n            \"onecyclelr\",\r\n            {\r\n                \"max_lr\": 1e-3,\r\n                \"epochs\": epochs,\r\n                \"steps_per_epoch\": steps_per_epoch,\r\n            },\r\n            {\"interval\": \"step\"},\r\n        ),\r\n    )\r\n```\r\n\r\n```python\r\ntrainer = flash.Trainer(\r\n    max_epochs=epochs,\r\n    gpus=torch.cuda.device_count(),\r\n)\r\ntrainer.finetune(model, datamodule=datamodule, strategy=(\"freeze_unfreeze\", 5))\r\n```\r\n\r\n### Expected behaviour\r\n\r\nAfter specified number of epochs, layers get unfrozen and training continues.\r\n\r\n### Actual behaviour\r\nExpection is thrown:\r\n\r\n```python\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py\", line 202, in start_training\r\n    self._results = trainer.run_stage()\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\", line 1289, in run_stage\r\n    return self._run_train()\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\", line 1319, in _run_train\r\n    self.fit_loop.run()\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\", line 145, in run\r\n    self.advance(*args, **kwargs)\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py\", line 234, in advance\r\n    self.epoch_loop.run(data_fetcher)\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\", line 145, in run\r\n    self.advance(*args, **kwargs)\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py\", line 199, in advance\r\n    self.update_lr_schedulers(\"step\", update_plateau_schedulers=False)\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py\", line 441, in update_lr_schedulers\r\n    self._update_learning_rates(\r\n  File \"/<redacted>/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py\", line 505, in _update_learning_rates\r\n    lr_scheduler[\"scheduler\"].step()\r\n  File \"/<redacted>/lib/python3.8/site-packages/torch/optim/lr_scheduler.py\", line 154, in step\r\n    values = self.get_lr()\r\n  File \"/<redacted>/lib/python3.8/site-packages/torch/optim/lr_scheduler.py\", line 1597, in get_lr\r\n    computed_lr = self.anneal_func(group[phase['start_lr']], group[phase['end_lr']], pct)\r\nKeyError: 'max_lr\r\n```\r\n\r\nIt 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`.\r\n\r\n### Environment\r\n\r\n - OS (e.g., Linux): macOS\r\n - Python version: 3.8.12\r\n - PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7): 1.11.0 / 1.5.10 / 0.7.3\r\n - GPU models and configuration: 0 / 1 T4 (happens regardless of cuda)\r\n - Any other relevant information:\r\n\r\n### Additional context\r\n\r\nhttps://pytorch-lightning.slack.com/archives/CRBLFHY79/p1651218144224359\r\n\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nfrom functools import partial\nfrom typing import Iterable, Optional, Tuple, Union\n\nfrom pytorch_lightning import LightningModule\nfrom pytorch_lightning.callbacks import BaseFinetuning\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.nn import Module\nfrom torch.optim import Optimizer\n\nfrom flash.core.registry import FlashRegistry\n\nif not os.environ.get(\"READTHEDOCS\", False):\n    from pytorch_lightning.utilities.enums import LightningEnum\nelse:\n    # ReadTheDocs mocks the `LightningEnum` import to be a regular type, so we replace it with a plain Enum here.\n    from enum import Enum\n\n    LightningEnum = Enum\n\n\nclass FinetuningStrategies(LightningEnum):\n    \"\"\"The ``FinetuningStrategies`` enum contains the keys that are used internally by the ``FlashBaseFinetuning``\n    when choosing the strategy to perform.\"\"\"\n\n    NO_FREEZE = \"no_freeze\"\n    FREEZE = \"freeze\"\n    FREEZE_UNFREEZE = \"freeze_unfreeze\"\n    UNFREEZE_MILESTONES = \"unfreeze_milestones\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass FlashBaseFinetuning(BaseFinetuning):\n    \"\"\"FlashBaseFinetuning can be used to create a custom Flash Finetuning Callback.\"\"\"\n\n    def __init__(\n        self,\n        strategy_key: Union[str, FinetuningStrategies],\n        strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = None,\n        train_bn: bool = True,\n    ):\n        \"\"\"\n        Args:\n            strategy_key: The finetuning strategy to be used. See :meth:`~flash.core.trainer.Trainer.finetune`\n                for the available strategies.\n            strategy_metadata: Data that accompanies certain finetuning strategies like epoch number or number of\n                layers.\n            train_bn: Whether to train Batch Norm layer\n        \"\"\"\n        super().__init__()\n\n        self.strategy: FinetuningStrategies = strategy_key\n        self.strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = strategy_metadata\n        self.train_bn: bool = train_bn\n\n        if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE and not isinstance(self.strategy_metadata, int):\n            raise MisconfigurationException(\n                \"The `freeze_unfreeze` strategy requires an integer denoting the epoch number to unfreeze at. Example: \"\n                \"`strategy=('freeze_unfreeze', 7)`\"\n            )\n        if self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES and not (\n            isinstance(self.strategy_metadata, Tuple)\n            and isinstance(self.strategy_metadata[0], Tuple)\n            and isinstance(self.strategy_metadata[1], int)\n            and isinstance(self.strategy_metadata[0][0], int)\n            and isinstance(self.strategy_metadata[0][1], int)\n        ):\n            raise MisconfigurationException(\n                \"The `unfreeze_milestones` strategy requires the format Tuple[Tuple[int, int], int]. Example: \"\n                \"`strategy=('unfreeze_milestones', ((5, 10), 15))`\"\n            )\n\n    def _get_modules_to_freeze(self, pl_module: LightningModule) -> Union[Module, Iterable[Union[Module, Iterable]]]:\n        modules_to_freeze = getattr(pl_module, \"modules_to_freeze\", None)\n        if modules_to_freeze is None:\n            raise AttributeError(\n                \"LightningModule missing instance method 'modules_to_freeze'.\"\n                \"Please, implement the method which returns NoneType or a Module or an Iterable of Modules.\"\n            )\n        return modules_to_freeze()\n\n    def freeze_before_training(self, pl_module: Union[Module, Iterable[Union[Module, Iterable]]]) -> None:\n        if self.strategy != FinetuningStrategies.NO_FREEZE:\n            modules = self._get_modules_to_freeze(pl_module=pl_module)\n            if modules is not None:\n                if isinstance(modules, Module):\n                    modules = [modules]\n                self.freeze(modules=modules, train_bn=self.train_bn)\n\n    def _freeze_unfreeze_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n        strategy_metadata: int,\n    ):\n        unfreeze_epoch: int = strategy_metadata\n        if epoch != unfreeze_epoch:\n            return\n\n        modules = self._get_modules_to_freeze(pl_module=pl_module)\n        if modules is not None:\n            self.unfreeze_and_add_param_group(\n                modules=modules,\n                optimizer=optimizer,\n                train_bn=self.train_bn,\n            )\n\n    def _unfreeze_milestones_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n        strategy_metadata: Tuple[Tuple[int, int], int],\n    ):\n        unfreeze_milestones: Tuple[int, int] = strategy_metadata[0]\n        num_layers: int = strategy_metadata[1]\n\n        modules = self._get_modules_to_freeze(pl_module=pl_module)\n        if modules is not None:\n            if epoch == unfreeze_milestones[0]:\n                # unfreeze num_layers last layers\n\n                backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[-num_layers:]\n                self.unfreeze_and_add_param_group(\n                    modules=backbone_modules,\n                    optimizer=optimizer,\n                    train_bn=self.train_bn,\n                )\n            elif epoch == unfreeze_milestones[1]:\n                # unfreeze remaining layers\n                backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[:-num_layers]\n                self.unfreeze_and_add_param_group(\n                    modules=backbone_modules,\n                    optimizer=optimizer,\n                    train_bn=self.train_bn,\n                )\n\n    def finetune_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n    ):\n        if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE:\n            self._freeze_unfreeze_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)\n        elif self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES:\n            self._unfreeze_milestones_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)\n\n\n_FINETUNING_STRATEGIES_REGISTRY = FlashRegistry(\"finetuning_strategies\")\n\nfor strategy in FinetuningStrategies:\n    _FINETUNING_STRATEGIES_REGISTRY(\n        name=strategy.value,\n        fn=partial(FlashBaseFinetuning, strategy_key=strategy),\n    )\n\n\nclass NoFreeze(FlashBaseFinetuning):\n    def __init__(self, train_bn: bool = True):\n        super().__init__(FinetuningStrategies.NO_FREEZE, train_bn)\n\n\nclass Freeze(FlashBaseFinetuning):\n    def __init__(self, train_bn: bool = True):\n        super().__init__(FinetuningStrategies.FREEZE, train_bn)\n\n\nclass FreezeUnfreeze(FlashBaseFinetuning):\n    def __init__(\n        self,\n        strategy_metadata: int,\n        train_bn: bool = True,\n    ):\n        super().__init__(FinetuningStrategies.FREEZE_UNFREEZE, strategy_metadata, train_bn)\n\n\nclass UnfreezeMilestones(FlashBaseFinetuning):\n    def __init__(\n        self,\n        strategy_metadata: Tuple[Tuple[int, int], int],\n        train_bn: bool = True,\n    ):\n        super().__init__(FinetuningStrategies.UNFREEZE_MILESTONES, strategy_metadata, train_bn)\n", "path": "flash/core/finetuning.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nfrom functools import partial\nfrom typing import Iterable, Optional, Tuple, Union\n\nfrom pytorch_lightning import LightningModule\nfrom pytorch_lightning.callbacks import BaseFinetuning\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.nn import Module\nfrom torch.optim import Optimizer\n\nfrom flash.core.registry import FlashRegistry\n\nif not os.environ.get(\"READTHEDOCS\", False):\n    from pytorch_lightning.utilities.enums import LightningEnum\nelse:\n    # ReadTheDocs mocks the `LightningEnum` import to be a regular type, so we replace it with a plain Enum here.\n    from enum import Enum\n\n    LightningEnum = Enum\n\n\nclass FinetuningStrategies(LightningEnum):\n    \"\"\"The ``FinetuningStrategies`` enum contains the keys that are used internally by the ``FlashBaseFinetuning``\n    when choosing the strategy to perform.\"\"\"\n\n    NO_FREEZE = \"no_freeze\"\n    FREEZE = \"freeze\"\n    FREEZE_UNFREEZE = \"freeze_unfreeze\"\n    UNFREEZE_MILESTONES = \"unfreeze_milestones\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass FlashBaseFinetuning(BaseFinetuning):\n    \"\"\"FlashBaseFinetuning can be used to create a custom Flash Finetuning Callback.\"\"\"\n\n    def __init__(\n        self,\n        strategy_key: Union[str, FinetuningStrategies],\n        strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = None,\n        train_bn: bool = True,\n    ):\n        \"\"\"\n        Args:\n            strategy_key: The finetuning strategy to be used. See :meth:`~flash.core.trainer.Trainer.finetune`\n                for the available strategies.\n            strategy_metadata: Data that accompanies certain finetuning strategies like epoch number or number of\n                layers.\n            train_bn: Whether to train Batch Norm layer\n        \"\"\"\n        super().__init__()\n\n        self.strategy: FinetuningStrategies = strategy_key\n        self.strategy_metadata: Optional[Union[int, Tuple[Tuple[int, int], int]]] = strategy_metadata\n        self.train_bn: bool = train_bn\n\n        if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE and not isinstance(self.strategy_metadata, int):\n            raise MisconfigurationException(\n                \"The `freeze_unfreeze` strategy requires an integer denoting the epoch number to unfreeze at. Example: \"\n                \"`strategy=('freeze_unfreeze', 7)`\"\n            )\n        if self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES and not (\n            isinstance(self.strategy_metadata, Tuple)\n            and isinstance(self.strategy_metadata[0], Tuple)\n            and isinstance(self.strategy_metadata[1], int)\n            and isinstance(self.strategy_metadata[0][0], int)\n            and isinstance(self.strategy_metadata[0][1], int)\n        ):\n            raise MisconfigurationException(\n                \"The `unfreeze_milestones` strategy requires the format Tuple[Tuple[int, int], int]. Example: \"\n                \"`strategy=('unfreeze_milestones', ((5, 10), 15))`\"\n            )\n\n    def _get_modules_to_freeze(self, pl_module: LightningModule) -> Union[Module, Iterable[Union[Module, Iterable]]]:\n        modules_to_freeze = getattr(pl_module, \"modules_to_freeze\", None)\n        if modules_to_freeze is None:\n            raise AttributeError(\n                \"LightningModule missing instance method 'modules_to_freeze'.\"\n                \"Please, implement the method which returns NoneType or a Module or an Iterable of Modules.\"\n            )\n        return modules_to_freeze()\n\n    def freeze_before_training(self, pl_module: Union[Module, Iterable[Union[Module, Iterable]]]) -> None:\n        if self.strategy != FinetuningStrategies.NO_FREEZE:\n            modules = self._get_modules_to_freeze(pl_module=pl_module)\n            if modules is not None:\n                if isinstance(modules, Module):\n                    modules = [modules]\n                self.freeze(modules=modules, train_bn=self.train_bn)\n\n    def unfreeze_and_extend_param_group(\n        self,\n        modules: Union[Module, Iterable[Union[Module, Iterable]]],\n        optimizer: Optimizer,\n        train_bn: bool = True,\n    ) -> None:\n        self.make_trainable(modules)\n\n        params = self.filter_params(modules, train_bn=train_bn, requires_grad=True)\n        params = self.filter_on_optimizer(optimizer, params)\n        if params:\n            optimizer.param_groups[0][\"params\"].extend(params)\n\n    def _freeze_unfreeze_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n        strategy_metadata: int,\n    ):\n        unfreeze_epoch: int = strategy_metadata\n        if epoch != unfreeze_epoch:\n            return\n\n        modules = self._get_modules_to_freeze(pl_module=pl_module)\n        if modules is not None:\n            self.unfreeze_and_extend_param_group(\n                modules=modules,\n                optimizer=optimizer,\n                train_bn=self.train_bn,\n            )\n\n    def _unfreeze_milestones_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n        strategy_metadata: Tuple[Tuple[int, int], int],\n    ):\n        unfreeze_milestones: Tuple[int, int] = strategy_metadata[0]\n        num_layers: int = strategy_metadata[1]\n\n        modules = self._get_modules_to_freeze(pl_module=pl_module)\n        if modules is not None:\n            if epoch == unfreeze_milestones[0]:\n                # unfreeze num_layers last layers\n\n                backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[-num_layers:]\n                self.unfreeze_and_extend_param_group(\n                    modules=backbone_modules,\n                    optimizer=optimizer,\n                    train_bn=self.train_bn,\n                )\n            elif epoch == unfreeze_milestones[1]:\n                # unfreeze remaining layers\n                backbone_modules = BaseFinetuning.flatten_modules(modules=modules)[:-num_layers]\n                self.unfreeze_and_extend_param_group(\n                    modules=backbone_modules,\n                    optimizer=optimizer,\n                    train_bn=self.train_bn,\n                )\n\n    def finetune_function(\n        self,\n        pl_module: Union[Module, Iterable[Union[Module, Iterable]]],\n        epoch: int,\n        optimizer: Optimizer,\n        opt_idx: int,\n    ):\n        if self.strategy == FinetuningStrategies.FREEZE_UNFREEZE:\n            self._freeze_unfreeze_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)\n        elif self.strategy == FinetuningStrategies.UNFREEZE_MILESTONES:\n            self._unfreeze_milestones_function(pl_module, epoch, optimizer, opt_idx, self.strategy_metadata)\n\n\n_FINETUNING_STRATEGIES_REGISTRY = FlashRegistry(\"finetuning_strategies\")\n\nfor strategy in FinetuningStrategies:\n    _FINETUNING_STRATEGIES_REGISTRY(\n        name=strategy.value,\n        fn=partial(FlashBaseFinetuning, strategy_key=strategy),\n    )\n\n\nclass NoFreeze(FlashBaseFinetuning):\n    def __init__(self, train_bn: bool = True):\n        super().__init__(FinetuningStrategies.NO_FREEZE, train_bn)\n\n\nclass Freeze(FlashBaseFinetuning):\n    def __init__(self, train_bn: bool = True):\n        super().__init__(FinetuningStrategies.FREEZE, train_bn)\n\n\nclass FreezeUnfreeze(FlashBaseFinetuning):\n    def __init__(\n        self,\n        strategy_metadata: int,\n        train_bn: bool = True,\n    ):\n        super().__init__(FinetuningStrategies.FREEZE_UNFREEZE, strategy_metadata, train_bn)\n\n\nclass UnfreezeMilestones(FlashBaseFinetuning):\n    def __init__(\n        self,\n        strategy_metadata: Tuple[Tuple[int, int], int],\n        train_bn: bool = True,\n    ):\n        super().__init__(FinetuningStrategies.UNFREEZE_MILESTONES, strategy_metadata, train_bn)\n", "path": "flash/core/finetuning.py"}]}