From b3e9dff32d842431b067b0ab83e508ffe3262968 Mon Sep 17 00:00:00 2001 From: Jirka Borovec Date: Wed, 6 Oct 2021 10:39:36 +0200 Subject: [PATCH] rename callback FineTune arg `round` (#9711) * rename CB Tune arg round Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- pytorch_lightning/callbacks/finetuning.py | 113 +++++++++----------- tests/callbacks/test_finetuning_callback.py | 4 +- 2 files changed, 50 insertions(+), 67 deletions(-) diff --git a/pytorch_lightning/callbacks/finetuning.py b/pytorch_lightning/callbacks/finetuning.py index 5a63a6e92d..f6cc7e0920 100644 --- a/pytorch_lightning/callbacks/finetuning.py +++ b/pytorch_lightning/callbacks/finetuning.py @@ -53,32 +53,29 @@ class BaseFinetuning(Callback): Example:: - class MyModel(LightningModule) - - ... - - def configure_optimizer(self): - # Make sure to filter the parameters based on `requires_grad` - return Adam(filter(lambda p: p.requires_grad, self.parameters)) - - class FeatureExtractorFreezeUnfreeze(BaseFinetuning): - - def __init__(self, unfreeze_at_epoch=10): - self._unfreeze_at_epoch = unfreeze_at_epoch - - def freeze_before_training(self, pl_module): - # freeze any module you want - # Here, we are freezing `feature_extractor` - self.freeze(pl_module.feature_extractor) - - def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx): - # When `current_epoch` is 10, feature_extractor will start training. - if current_epoch == self._unfreeze_at_epoch: - self.unfreeze_and_add_param_group( - modules=pl_module.feature_extractor, - optimizer=optimizer, - train_bn=True, - ) + >>> from torch.optim import Adam + >>> class MyModel(pl.LightningModule): + ... def configure_optimizer(self): + ... # Make sure to filter the parameters based on `requires_grad` + ... return Adam(filter(lambda p: p.requires_grad, self.parameters)) + ... + >>> class FeatureExtractorFreezeUnfreeze(BaseFinetuning): + ... def __init__(self, unfreeze_at_epoch=10): + ... self._unfreeze_at_epoch = unfreeze_at_epoch + ... + ... def freeze_before_training(self, pl_module): + ... # freeze any module you want + ... # Here, we are freezing `feature_extractor` + ... self.freeze(pl_module.feature_extractor) + ... + ... def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx): + ... # When `current_epoch` is 10, feature_extractor will start training. + ... if current_epoch == self._unfreeze_at_epoch: + ... self.unfreeze_and_add_param_group( + ... modules=pl_module.feature_extractor, + ... optimizer=optimizer, + ... train_bn=True, + ... ) """ def __init__(self): @@ -208,7 +205,7 @@ class BaseFinetuning(Callback): if removed_params: rank_zero_warn( - "The provided params to be freezed already exist within another group of this optimizer." + "The provided params to be frozen already exist within another group of this optimizer." " Those parameters will be skipped.\n" "HINT: Did you init your optimizer in `configure_optimizer` as such:\n" f" {type(optimizer)}(filter(lambda p: p.requires_grad, self.parameters()), ...) ", @@ -227,22 +224,14 @@ class BaseFinetuning(Callback): """Unfreezes a module and adds its parameters to an optimizer. Args: - modules: A module or iterable of modules to unfreeze. Their parameters will be added to an optimizer as a new param group. - optimizer: The provided optimizer will receive new parameters and will add them to `add_param_group` - lr: Learning rate for the new param group. - initial_denom_lr: If no lr is provided, the learning from the first param group will be used - and divided by initial_denom_lr. - + and divided by `initial_denom_lr`. train_bn: Whether to train the BatchNormalization layers. - - Returns: - None """ BaseFinetuning.make_trainable(modules) params_lr = optimizer.param_groups[0]["lr"] if lr is None else float(lr) @@ -252,7 +241,7 @@ class BaseFinetuning(Callback): if params: optimizer.add_param_group({"params": params, "lr": params_lr / denom_lr}) - def on_before_accelerator_backend_setup(self, trainer, pl_module): + def on_before_accelerator_backend_setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"): self.freeze_before_training(pl_module) @staticmethod @@ -283,7 +272,7 @@ class BaseFinetuning(Callback): self.__apply_mapping_to_param_groups(current_param_groups[num_param_groups:], mapping) ) - def on_train_epoch_start(self, trainer, pl_module): + def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: """Called when the epoch begins.""" # import is here to avoid circular imports from pytorch_lightning.loops.utilities import _get_active_optimizers @@ -294,45 +283,37 @@ class BaseFinetuning(Callback): current_param_groups = optimizer.param_groups self._store(pl_module, opt_idx, num_param_groups, current_param_groups) - def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int): + def finetune_function( + self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int + ) -> None: """Override to add your unfreeze logic.""" raise NotImplementedError - def freeze_before_training(self, pl_module: "pl.LightningModule"): + def freeze_before_training(self, pl_module: "pl.LightningModule") -> None: """Override to add your freeze logic.""" raise NotImplementedError class BackboneFinetuning(BaseFinetuning): - r""" + r"""Finetune a backbone model based on a learning rate user-defined scheduling. - Finetune a backbone model based on a learning rate user-defined scheduling. When the backbone learning rate reaches the current model learning rate and ``should_align`` is set to True, it will align with it for the rest of the training. Args: - unfreeze_backbone_at_epoch: Epoch at which the backbone will be unfreezed. - lambda_func: Scheduling function for increasing backbone learning rate. - backbone_initial_ratio_lr: Used to scale down the backbone learning rate compared to rest of model - backbone_initial_lr: Optional, Inital learning rate for the backbone. - By default, we will use current_learning / backbone_initial_ratio_lr - + By default, we will use ``current_learning / backbone_initial_ratio_lr`` should_align: Wheter to align with current learning rate when backbone learning reaches it. - initial_denom_lr: When unfreezing the backbone, the intial learning rate will - current_learning_rate / initial_denom_lr. - + ``current_learning_rate / initial_denom_lr``. train_bn: Wheter to make Batch Normalization trainable. - verbose: Display current learning rate for model and backbone - - round: Precision for displaying learning rate + rounding: Precision for displaying learning rate Example:: @@ -354,8 +335,8 @@ class BackboneFinetuning(BaseFinetuning): initial_denom_lr: float = 10.0, train_bn: bool = True, verbose: bool = False, - round: int = 12, - ): + rounding: int = 12, + ) -> None: super().__init__() self.unfreeze_backbone_at_epoch: int = unfreeze_backbone_at_epoch @@ -366,12 +347,12 @@ class BackboneFinetuning(BaseFinetuning): self.initial_denom_lr: float = initial_denom_lr self.train_bn: bool = train_bn self.verbose: bool = verbose - self.round: int = round + self.rounding: int = rounding self.previous_backbone_lr: Optional[float] = None def on_save_checkpoint( self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any] - ) -> Dict[int, Any]: + ) -> Dict[str, Any]: return { "internal_optimizer_metadata": self._internal_optimizer_metadata, "previous_backbone_lr": self.previous_backbone_lr, @@ -383,7 +364,7 @@ class BackboneFinetuning(BaseFinetuning): self.previous_backbone_lr = callback_state["previous_backbone_lr"] super().on_load_checkpoint(trainer, pl_module, callback_state["internal_optimizer_metadata"]) - def on_fit_start(self, trainer, pl_module): + def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: """ Raises: MisconfigurationException: @@ -393,10 +374,12 @@ class BackboneFinetuning(BaseFinetuning): return super().on_fit_start(trainer, pl_module) raise MisconfigurationException("The LightningModule should have a nn.Module `backbone` attribute") - def freeze_before_training(self, pl_module: "pl.LightningModule"): + def freeze_before_training(self, pl_module: "pl.LightningModule") -> None: self.freeze(pl_module.backbone) - def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int): + def finetune_function( + self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int + ) -> None: """Called when the epoch begins.""" if epoch == self.unfreeze_backbone_at_epoch: current_lr = optimizer.param_groups[0]["lr"] @@ -415,8 +398,8 @@ class BackboneFinetuning(BaseFinetuning): ) if self.verbose: log.info( - f"Current lr: {round(current_lr, self.round)}, " - f"Backbone lr: {round(initial_backbone_lr, self.round)}" + f"Current lr: {round(current_lr, self.rounding)}, " + f"Backbone lr: {round(initial_backbone_lr, self.rounding)}" ) elif epoch > self.unfreeze_backbone_at_epoch: @@ -431,6 +414,6 @@ class BackboneFinetuning(BaseFinetuning): self.previous_backbone_lr = next_current_backbone_lr if self.verbose: log.info( - f"Current lr: {round(current_lr, self.round)}, " - f"Backbone lr: {round(next_current_backbone_lr, self.round)}" + f"Current lr: {round(current_lr, self.rounding)}, " + f"Backbone lr: {round(next_current_backbone_lr, self.rounding)}" ) diff --git a/tests/callbacks/test_finetuning_callback.py b/tests/callbacks/test_finetuning_callback.py index 31b8b0e160..5a0b6e33c5 100644 --- a/tests/callbacks/test_finetuning_callback.py +++ b/tests/callbacks/test_finetuning_callback.py @@ -186,7 +186,7 @@ def test_unfreeze_and_add_param_group_function(tmpdir): model = FreezeModel() optimizer = SGD(model.backbone[0].parameters(), lr=0.01) - with pytest.warns(UserWarning, match="The provided params to be freezed already"): + with pytest.warns(UserWarning, match="The provided params to be frozen already"): BaseFinetuning.unfreeze_and_add_param_group(model.backbone[0], optimizer=optimizer) assert optimizer.param_groups[0]["lr"] == 0.01 @@ -197,7 +197,7 @@ def test_unfreeze_and_add_param_group_function(tmpdir): assert torch.equal(optimizer.param_groups[1]["params"][0], model.backbone[1].weight) assert model.backbone[1].weight.requires_grad - with pytest.warns(UserWarning, match="The provided params to be freezed already"): + with pytest.warns(UserWarning, match="The provided params to be frozen already"): BaseFinetuning.unfreeze_and_add_param_group(model, optimizer=optimizer, lr=100, train_bn=False) assert len(optimizer.param_groups) == 3 assert optimizer.param_groups[2]["lr"] == 100