From c75c108384c9992f5b28010491fe0a1191312a3b Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Thu, 5 May 2022 17:30:41 +0100 Subject: [PATCH 1/3] Fix unfreeze strategies with onecyclelr and reduced lr --- CHANGELOG.md | 4 ++++ flash/core/finetuning.py | 19 +++++++++++++++--- tests/core/test_finetuning.py | 37 +++++++++++++++++++++++++++++------ 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 908b964a0c..b25a4de30b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed a bug where a loaded `TabularClassifier` or `TabularRegressor` checkpoint could not be served ([#1324](https://github.com/PyTorchLightning/lightning-flash/pull/1324)) +- Fixed a bug where the `freeze_unfreeze` and `unfreeze_milestones` finetuning strategies could not be used in tandem with a `onecyclelr` LR scheduler + +- Fixed a bug where the backbone learning rate would be divided by 10 when unfrozen if using the `freeze_unfreeze` or `unfreeze_milestones` strategies + ## [0.7.4] - 2022-04-27 ### Fixed diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index 4038f1e9ca..f6b39c928e 100644 --- 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, diff --git a/tests/core/test_finetuning.py b/tests/core/test_finetuning.py index 88f5b9851d..ea370c758f 100644 --- a/tests/core/test_finetuning.py +++ b/tests/core/test_finetuning.py @@ -152,20 +152,45 @@ def test_finetuning_with_none_return_type(strategy): @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() From 3307c937db7455c8e6b8abe51d23432183e55587 Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Thu, 5 May 2022 17:34:23 +0100 Subject: [PATCH 2/3] Update CHANGELOG.md --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b25a4de30b..7c2f441433 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,9 +24,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed a bug where a loaded `TabularClassifier` or `TabularRegressor` checkpoint could not be served ([#1324](https://github.com/PyTorchLightning/lightning-flash/pull/1324)) -- Fixed a bug where the `freeze_unfreeze` and `unfreeze_milestones` finetuning strategies could not be used in tandem with a `onecyclelr` LR scheduler +- Fixed a bug where the `freeze_unfreeze` and `unfreeze_milestones` finetuning strategies could not be used in tandem with a `onecyclelr` LR scheduler ([#1329](https://github.com/PyTorchLightning/lightning-flash/pull/1329)) -- Fixed a bug where the backbone learning rate would be divided by 10 when unfrozen if using the `freeze_unfreeze` or `unfreeze_milestones` strategies +- Fixed a bug where the backbone learning rate would be divided by 10 when unfrozen if using the `freeze_unfreeze` or `unfreeze_milestones` strategies ([#1329](https://github.com/PyTorchLightning/lightning-flash/pull/1329)) ## [0.7.4] - 2022-04-27 From b77ecf91fc0644c38a34522d7dadcbd9cd854451 Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Fri, 6 May 2022 14:48:07 +0100 Subject: [PATCH 3/3] Update docs --- docs/source/general/finetuning.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/general/finetuning.rst b/docs/source/general/finetuning.rst index fdd31146d8..42cb873e29 100644 --- a/docs/source/general/finetuning.rst +++ b/docs/source/general/finetuning.rst @@ -228,7 +228,7 @@ For even more customization, create your own finetuning callback. Learn more abo # When ``current_epoch`` is 5, backbone will start to be trained. if current_epoch == self._unfreeze_epoch: - self.unfreeze_and_add_param_group( + self.unfreeze_and_extend_param_group( pl_module.backbone, optimizer, )