Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#21,531 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

needs triage refactor
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Outline & Motivation

The [`BackboneFinetuning`](https://github.com/Lightning-AI/pytorch-lightning/blob/215d7d9737fc20168b188b71003f355357904d53/src/lightning/pytorch/callbacks/finetuning.py#L341) 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:
```python
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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/lightning/pytorch/callbacks/finetuning.py at BackboneFinetuning.freeze_before_training() and review how train_bn is handled in the frozen and unfreezing phases. Check the callback’s parameter handling and existing callback tests; done when both phase-specific settings work, the legacy parameter emits the requested deprecation warning, and current defaults remain compatible.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.