fix(training): restore efficientnet-lite0 fine-tuning quality lost in timm swap (#16)

Switching the efficientnet-lite0 backbone to timm silently dropped three
model-level defaults that the previous package applied, regressing FixMatch
fine-tuning AUROC (~0.96 -> ~0.81 on the internal benchmark):

- Classifier head over-scaling (main cause): timm initialises the head with
  TF-EfficientNet's 1/sqrt(fan_in+fan_out) scale tuned for 1000 classes, which
  is ~num_classes too large for the 2-class head (weight std ~0.4 vs PyTorch's
  ~0.016). The overconfident fresh head makes ~all unlabeled images clear the
  FixMatch p_cutoff from step 1 with random pseudo-labels, poisoning the
  backbone. Reset the classifier to PyTorch's default init after create_model.
- Unapplied bn_momentum: cfg.bn_momentum (= 1 - ema_m, ~0.01) was computed and
  validated but never set on the model, so BatchNorm ran at timm's 0.1 default
  (~10x too fast for the small batch size). Apply it to both models' BatchNorm.
- Unapplied seed: cfg.seed was defined and validated but set_seeds was never
  called, so training was not reproducible. Seed all RNGs after validation.

Add a regression test asserting the classifier head weight std stays near
PyTorch's Linear bound.

Author: gomezzz

anomaly_match/pipeline/session.py

@@ -36,6 +36,7 @@
 from anomaly_match.utils.get_optimizer import get_optimizer
 from anomaly_match.utils.print_cfg import print_cfg
 from anomaly_match.utils.set_log_level import set_log_level
+from anomaly_match.utils.set_seeds import set_seeds
 from anomaly_match.utils.validate_config import validate_config
 
 
@@ -83,6 +84,11 @@ def __init__(self, cfg):
         # Validate the config
         validate_config(cfg)
 
+        # Seed all RNGs before datasets and model are built so that the train/test
+        # split, weight initialisation and augmentation sampling are reproducible.
+        # cfg.seed was defined and validated but never actually applied.
+        set_seeds(int(cfg.seed))
+
         self.cfg = cfg
         self.cached_image_normalisation_enum = cfg.normalisation.normalisation_method
         self.out = None  # Initialize out attribute to None
@@ -116,6 +122,15 @@ def _init_model(self):
             session_tracker=self.session_tracker,
         )
 
+        # Apply the configured BatchNorm momentum (cfg.bn_momentum = 1 - ema_m, ~0.01).
+        # It was computed and validated but never set on the model, so BatchNorm ran at
+        # timm's 0.1 default (~10x too fast for our small batch size), destabilising the
+        # running statistics during fine-tuning.
+        for submodel in (self.model.train_model, self.model.eval_model):
+            for module in submodel.modules():
+                if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+                    module.momentum = self.cfg.bn_momentum
+
         # get optimizer, ADAM and SGD are supported.
         optimizer = get_optimizer(
             self.model.train_model,

anomaly_match/utils/get_net_builder.py

@@ -95,6 +95,30 @@ def _resolve_timm_name(net_name, pretrained):
     return timm_base, False
 
 
+def _reset_classifier_head(model):
+    """Reset the classifier head to PyTorch's default Linear initialisation.
+
+    timm initialises EfficientNet classifiers with TF-EfficientNet's
+    ``1/sqrt(fan_in + fan_out)`` scale, which is tuned for the 1000-class ImageNet
+    head. For AnomalyMatch's 2-class head this produces weights that are roughly
+    ``num_classes`` times too large (std ~0.4 vs PyTorch's ~0.016). Such an
+    overconfident fresh head makes almost every unlabeled image clear the FixMatch
+    ``p_cutoff`` from the very first step with essentially random pseudo-labels,
+    poisoning the backbone during fine-tuning. Resetting to PyTorch's default
+    kaiming-uniform init restores the pre-timm fine-tuning quality.
+
+    Args:
+        model: A timm model exposing ``get_classifier()``.
+
+    Returns:
+        The same model, with its classifier head re-initialised in place.
+    """
+    classifier = model.get_classifier()
+    if isinstance(classifier, nn.Linear):
+        classifier.reset_parameters()
+    return model
+
+
 def get_net_builder(net_name, pretrained=False, in_channels=3):
     """Create a neural network builder function for the specified architecture.
 
@@ -136,7 +160,7 @@ def build_model(
         effective_pretrained = pretrained if pretrained is not None else _pretrained
         if effective_pretrained:
             try:
-                return timm.create_model(
+                model = timm.create_model(
                     _timm_name,
                     pretrained=True,
                     num_classes=num_classes,
@@ -148,17 +172,22 @@ def build_model(
                     f"Bundled pretrained weights not available (clone with git-lfs to avoid "
                     f"re-downloading). Downloading {_timm_name} from HuggingFace."
                 )
-                return timm.create_model(
+                model = timm.create_model(
                     _timm_name,
                     pretrained=True,
                     num_classes=num_classes,
                     in_chans=in_channels,
                 )
-        return timm.create_model(
-            _timm_name,
-            pretrained=False,
-            num_classes=num_classes,
-            in_chans=in_channels,
-        )
+        else:
+            model = timm.create_model(
+                _timm_name,
+                pretrained=False,
+                num_classes=num_classes,
+                in_chans=in_channels,
+            )
+        # timm's fresh classifier head is scaled for the 1000-class ImageNet head and
+        # is ~num_classes too large for AnomalyMatch's 2-class head; reset it to
+        # PyTorch's default init to avoid poisoning FixMatch fine-tuning.
+        return _reset_classifier_head(model)
 
     return build_model

tests/unit/test_net_builder.py

@@ -0,0 +1,24 @@
+#   Copyright (c) European Space Agency, 2025.
+#
+#   This file is subject to the terms and conditions defined in file 'LICENCE.txt', which
+#   is part of this source code package. No part of the package, including
+#   this file, may be copied, modified, propagated, or distributed except according to
+#   the terms contained in the file 'LICENCE.txt'.
+import torch.nn as nn
+
+from anomaly_match.utils.get_net_builder import get_net_builder
+
+
+def test_efficientnet_classifier_head_init_is_pytorch_scale():
+    """The 2-class classifier head must use PyTorch's default init, not timm's.
+
+    timm scales the fresh EfficientNet head for the 1000-class ImageNet head, which is
+    far too large for AnomalyMatch's 2-class head and poisons FixMatch fine-tuning.
+    After the head reset the weight std must sit near PyTorch's Linear bound.
+    """
+    model = get_net_builder("efficientnet-lite0", pretrained=False)(num_classes=2, in_channels=3)
+    classifier = model.get_classifier()
+    assert isinstance(classifier, nn.Linear)
+    weight = classifier.weight
+    pytorch_bound = 1.0 / (weight.shape[1] ** 0.5)
+    assert weight.std().item() < 2 * pytorch_bound