Skip to content

Fix efficientnet-lite0 fine-tuning regression from the timm backbone swap - #15

Closed
gomezzz wants to merge 1 commit into
esa:mainfrom
gomezzz:fix/efficientnet-lite0-training-regression
Closed

Fix efficientnet-lite0 fine-tuning regression from the timm backbone swap#15
gomezzz wants to merge 1 commit into
esa:mainfrom
gomezzz:fix/efficientnet-lite0-training-regression

Conversation

@gomezzz

@gomezzz gomezzz commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

Switching the efficientnet-lite0 backbone to timm (tf_efficientnet_lite0) silently dropped three model-level defaults that the previous efficientnet_lite_pytorch package applied. Together they regressed FixMatch fine-tuning quality substantially. This PR restores all three:

  • Classifier head over-scaling (main cause). timm initialises the head with TF-EfficientNet's 1/sqrt(fan_in+fan_out) scale, tuned for the 1000-class ImageNet head — roughly num_classes too large for AnomalyMatch's 2-class head (weight std ~0.4 vs PyTorch's ~0.016). The overconfident fresh head makes ~100% of unlabeled images clear the FixMatch p_cutoff (0.95) from the very first step with essentially random pseudo-labels, poisoning the backbone. Fix: reset the classifier to PyTorch's default init after timm.create_model(...).
  • cfg.bn_momentum never applied. It is computed (= 1 - ema_m, ~0.01) 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). Fix: apply it to both train_model and eval_model BatchNorm layers.
  • cfg.seed never applied. The seed was defined and validated but set_seeds was never called, so training was not reproducible. Fix: seed all RNGs after config validation, before datasets and model are built.

A regression test asserts the classifier head weight std stays near PyTorch's Linear bound.

Validation (internal benchmark, AUROC)

Class Before (timm regression) After fix
hourglass 0.81 0.94 – 0.96
piano 0.83 0.89 – 0.92

Test plan

  • ruff check + ruff format clean on changed files
  • New tests/unit/test_net_builder.py::test_efficientnet_classifier_head_init_is_pytorch_scale passes
  • tests/unit/test_fixmatch.py and tests/unit/test_import.py still pass
  • anomaly_match.pipeline.session imports cleanly

… 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.
@gomezzz

gomezzz commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator Author

Superseded by #16, which uses a direct branch on this repo (no fork) — this clears the fork-only coverage-comment CI restriction. Closing in favour of #16.

@gomezzz gomezzz closed this Jul 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant