freeze Method in BackboneFinetuning Missing train_bn Argument

[Hadisd]

Bug description

In BackboneFinetuning, which extends BaseFinetuning, the freeze_before_training method calls self.freeze(pl_module.backbone). However, the freeze method requires an additional train_bn argument to properly handle batch normalization layers.

Affected File
File: pytorch-lightning/src/lightning/pytorch/callbacks/finetuning.py
Class: BackboneFinetuning
Method: freeze_before_training

@override
def freeze_before_training(self, pl_module: "pl.LightningModule") -> None:
    self.freeze(pl_module.backbone)  # Missing `train_bn`

What version are you seeing the problem on?

v2.5

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[SkafteNicki]

The train_bn argument does not seem to be a required argument but an optional one

pytorch-lightning/src/lightning/pytorch/callbacks/finetuning.py

Lines 200 to 216 in e752cb5

	def freeze(modules: Union[Module, Iterable[Union[Module, Iterable]]], train_bn: bool = True) -> None:
	"""Freezes the parameters of the provided modules.
	Args:
	modules: A given module or an iterable of modules
	train_bn: If True, leave the BatchNorm layers in training mode
	Returns:
	None
	"""
	modules = BaseFinetuning.flatten_modules(modules)
	for mod in modules:
	if isinstance(mod, _BatchNorm) and train_bn:
	BaseFinetuning.make_trainable(mod)
	else:
	BaseFinetuning.freeze_module(mod)

It is set to True by default which will means that batch normalization layers by default is left trainable.

[SkafteNicki]

If more fine grained control is needed over the finetuning then the BackboneFinetuning.freeze_before_training method can be override by the user. But from experiment it seems to work as intended that batch normalization layers are still marked as trainable and other layers are frozen.
Closing issue.

[Hadisd]

Thank you. So, there's no need for the train_bn argument in BackboneFinetuning, since it behaves as if it's always set to True, regardless of its actual value.

Line 373 and 415
`class BackboneFinetuning(BaseFinetuning):
r"""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, Initial learning rate for the backbone.
        By default, we will use ``current_learning /  backbone_initial_ratio_lr``
    should_align: Whether to align with current learning rate when backbone learning
        reaches it.
    initial_denom_lr: When unfreezing the backbone, the initial learning rate will
        ``current_learning_rate /  initial_denom_lr``.
    train_bn: Whether to make Batch Normalization trainable.
    verbose: Display current learning rate for model and backbone
    rounding: Precision for displaying learning rate

Example::

    >>> from lightning.pytorch import Trainer
    >>> from lightning.pytorch.callbacks import BackboneFinetuning
    >>> multiplicative = lambda epoch: 1.5
    >>> backbone_finetuning = BackboneFinetuning(200, multiplicative)
    >>> trainer = Trainer(callbacks=[backbone_finetuning])

"""

def __init__(
    self,
    unfreeze_backbone_at_epoch: int = 10,
    lambda_func: Callable = multiplicative,
    backbone_initial_ratio_lr: float = 10e-2,
    backbone_initial_lr: Optional[float] = None,
    should_align: bool = True,
    initial_denom_lr: float = 10.0,
    train_bn: bool = True, # if you set it false it dosnt effect 
    verbose: bool = False,
    rounding: int = 12,
) -> None:
    super().__init__()

    self.unfreeze_backbone_at_epoch: int = unfreeze_backbone_at_epoch
    self.lambda_func: Callable = lambda_func
    self.backbone_initial_ratio_lr: float = backbone_initial_ratio_lr
    self.backbone_initial_lr: Optional[float] = backbone_initial_lr
    self.should_align: bool = should_align
    self.initial_denom_lr: float = initial_denom_lr
    self.train_bn: bool = train_bn
    self.verbose: bool = verbose
    self.rounding: int = rounding
    self.previous_backbone_lr: Optional[float] = None

@override
def state_dict(self) -> dict[str, Any]:
    return {
        "internal_optimizer_metadata": self._internal_optimizer_metadata,
        "previous_backbone_lr": self.previous_backbone_lr,
    }

@override
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    self.previous_backbone_lr = state_dict["previous_backbone_lr"]
    super().load_state_dict(state_dict)

@override
def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
    """
    Raises:
        MisconfigurationException:
            If LightningModule has no nn.Module `backbone` attribute.
    """
    if hasattr(pl_module, "backbone") and isinstance(pl_module.backbone, Module):
        return super().on_fit_start(trainer, pl_module)
    raise MisconfigurationException("The LightningModule should have a nn.Module `backbone` attribute")

@override
def freeze_before_training(self, pl_module: "pl.LightningModule") -> None:
    self.freeze(pl_module.backbone) # becasue it default value is always true.

`