Lightning-AI / Lightning-AI/pytorch-lightning
`BackboneFinetuning`: `train_bn` only applied during unfreezing phase
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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