From f06f034b98d5399ff1109c764d59a9c197d1a111 Mon Sep 17 00:00:00 2001 From: Pablo Gomez Date: Mon, 6 Jul 2026 20:41:24 +0000 Subject: [PATCH] fix(training): restore efficientnet-lite0 fine-tuning quality lost in timm swap 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. --- anomaly_match/pipeline/session.py | 15 +++++++++ anomaly_match/utils/get_net_builder.py | 45 +++++++++++++++++++++----- tests/unit/test_net_builder.py | 24 ++++++++++++++ 3 files changed, 76 insertions(+), 8 deletions(-) create mode 100644 tests/unit/test_net_builder.py diff --git a/anomaly_match/pipeline/session.py b/anomaly_match/pipeline/session.py index 77c60df..56217e8 100644 --- a/anomaly_match/pipeline/session.py +++ b/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, diff --git a/anomaly_match/utils/get_net_builder.py b/anomaly_match/utils/get_net_builder.py index 8a3f2a9..4862850 100644 --- a/anomaly_match/utils/get_net_builder.py +++ b/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 diff --git a/tests/unit/test_net_builder.py b/tests/unit/test_net_builder.py new file mode 100644 index 0000000..adde4f3 --- /dev/null +++ b/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