`BackboneFinetuning`: `train_bn` only applied during unfreezing phase

[patrontheo]

Outline & Motivation

The BackboneFinetuning callback exposes a train_bn parameter intended to control whether BatchNorm layers are trainable during backbone finetuning. However, the current implementation only applies this parameter during the unfreezing phase, not during the initial frozen phase.

In freeze_before_training(), the callback calls:

self.freeze(pl_module.backbone)

which uses the default train_bn=True. As a result, BatchNorm layers remain trainable during the frozen stage, regardless of the train_bn value passed to the callback.

This leads to a somewhat counter-intuitive behavior if train_bn=False:

• It does not freeze BN during the frozen phase.
• it freezes BN when the backbone is unfrozen.

So the meaning of the parameter becomes:
“Train BN while the backbone is frozen, and optionally freeze it once the backbone is unfrozen.”
This is not what the parameter name suggests, and is rarely the intended finetuning strategy.

Phase	train_bn=True	train_bn=False
Frozen phase	Backbone: frozen BN: trainable	Backbone: frozen BN: trainable
After unfreeze	Backbone: trainable BN: trainable	Backbone: trainable BN: frozen

Pitch

To keep current behavior available while making BN handling explicit and predictable:

• Deprecate the existing train_bn parameter.
• Introduce two new parameters:
  • train_bn_frozen_phase: controls whether BatchNorm layers are trainable while the backbone is frozen.
  • train_bn_unfrozen_phase: controls whether BatchNorm layers are trainable after the backbone is unfrozen.
• Set the default values to match the current behavior:
  • train_bn_frozen_phase=True
  • train_bn_unfrozen_phase=True
• Keep the old train_bn parameter for one deprecation cycle, mapping it internally to train_bn_unfrozen_phase=train_bn
• Emit a deprecation warning when train_bn is used, directing users to the new parameters.
• Remove train_bn in a future major release once the transition period is over.

Happy to discuss any other directions / improvements you have in mind.

Additional context

I’m happy to open a PR if the direction makes sense.

cc @lantiga @justusschock

[patrontheo]

I would like to add a comment on filter_params behavior and docstring.

The current implementation of filter_params is:

    @staticmethod
    def filter_params(
        modules: Union[Module, Iterable[Union[Module, Iterable]]], train_bn: bool = True, requires_grad: bool = True
    ) -> Generator:
        """Yields the `requires_grad` parameters of a given module or list of modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: Whether not to train the BatchNorm module
            requires_grad: Whether to create a generator for trainable or non-trainable parameters.
        Returns:
            Generator

        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and not train_bn:
                continue
            # recursion could yield duplicate parameters for parent modules w/ parameters so disabling it
            for param in mod.parameters(recurse=False):
                if param.requires_grad == requires_grad:
                    yield param

Docstring mismatch

The docstring says:

Yields the requires_grad parameters of a given module or list of modules.

This wording suggests that the function always yields parameters with requires_grad=True. However, the function actually has a requires_grad argument that allows selecting either:

• trainable parameters (requires_grad=True), or
• frozen parameters (requires_grad=False).

So the function is really filtering parameters based on their requires_grad value, not just yielding the trainable ones. A more accurate description would be something like:

“Yield parameters from the given modules whose requires_grad attribute matches the requires_grad argument, optionally skipping BatchNorm parameters when train_bn=False.”

Or

"Yield parameters from the given modules based on their requires_grad state, optionally skipping BatchNorm parameters when train_bn=False."

Interaction between train_bn and requires_grad=False

There is also a potentially confusing interaction between the two arguments:

if isinstance(mod, _BatchNorm) and not train_bn:
    continue

This condition skips all BatchNorm parameters when train_bn=False, regardless of their requires_grad state.

So if the function is called with:

filter_params(..., train_bn=False, requires_grad=False)

the intent would normally be:

Give me all frozen parameters.

But BatchNorm parameters are still excluded, even if their requires_grad=False.
This means the function is no longer acting as a pure “filter by requires_grad” utility, because the train_bn flag changes the output even when selecting frozen parameters.

This behavior may be surprising, since the name train_bn suggests it should only affect which parameters are considered trainable, not which frozen parameters are returned.

At minimum, this interaction could be clarified in the docstring, or the logic could be adjusted so that train_bn only affects the requires_grad=True case.