Skip to content

Commit ffa3c5a

Browse files
committed
add example for encoder freezing
1 parent 77fd490 commit ffa3c5a

File tree

1 file changed

+22
-0
lines changed

1 file changed

+22
-0
lines changed

docs/insights.rst

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,3 +117,25 @@ Example:
117117
118118
mask.shape, label.shape
119119
# (N, 4, H, W), (N, 4)
120+
121+
4. Freezing and unfreezing the encoder
122+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
123+
124+
Sometimes you may want to freeze the encoder during training, e.g. when using pretrained backbones and only fine-tuning the decoder and segmentation head.
125+
126+
All segmentation models in SMP provide two helper methods:
127+
128+
.. code-block:: python
129+
130+
model = smp.Unet("resnet34", classes=2)
131+
132+
# Freeze encoder: stops gradient updates and freezes normalization layer stats
133+
model.freeze_encoder()
134+
135+
# Unfreeze encoder: re-enables training for encoder parameters and normalization layers
136+
model.unfreeze_encoder()
137+
138+
.. important::
139+
- Freezing sets ``requires_grad = False`` for all encoder parameters.
140+
- Normalization layers that track running statistics (e.g., BatchNorm and InstanceNorm layers) are set to ``.eval()`` mode to prevent updates to ``running_mean`` and ``running_var``.
141+
- If you later call ``model.train()``, frozen encoders will remain frozen until you call ``unfreeze_encoder()``.

0 commit comments

Comments
 (0)