Add tests for encoder freezing

Author: ricber

tests/base/test_freeze_encoder.py

@@ -0,0 +1,54 @@
+import torch
+import segmentation_models_pytorch as smp
+
+
+def test_freeze_and_unfreeze_encoder():
+    model = smp.Unet(encoder_name="resnet18", encoder_weights=None)
+    model.train()
+    # Initially, encoder params should be trainable
+    assert all(p.requires_grad for p in model.encoder.parameters())
+    model.freeze_encoder()
+    # Check encoder params are frozen
+    assert all(not p.requires_grad for p in model.encoder.parameters())
+    # Check normalization layers are in eval mode
+    for m in model.encoder.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._NormBase):
+            assert not m.training
+    # Call train() and ensure encoder norm layers stay frozen
+    model.train()
+    for m in model.encoder.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._NormBase):
+            assert not m.training
+    model.unfreeze_encoder()
+    # Params should be trainable again
+    assert all(p.requires_grad for p in model.encoder.parameters())
+    # Norm layers should go back to training mode after unfreeze
+    for m in model.encoder.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._NormBase):
+            assert m.training
+    model.train()
+    # Norm layers should have the same training mode after train()
+    for m in model.encoder.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._NormBase):
+            assert m.training
+
+
+def test_freeze_encoder_stops_running_stats():
+    model = smp.Unet(encoder_name="resnet18", encoder_weights=None)
+    model.freeze_encoder()
+    model.train()  # overridden train, encoder should remain frozen
+    bn = None
+
+    for m in model.encoder.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._NormBase):
+            bn = m
+            break
+
+    assert bn is not None
+
+    orig_mean = bn.running_mean.clone()
+
+    x = torch.randn(2, 3, 64, 64)
+    _ = model(x)
+
+    torch.testing.assert_close(orig_mean, bn.running_mean)