Fix efficientnet-lite0 fine-tuning regression from the timm backbone swap - #15
Closed
gomezzz wants to merge 1 commit into
Closed
Fix efficientnet-lite0 fine-tuning regression from the timm backbone swap#15gomezzz wants to merge 1 commit into
gomezzz wants to merge 1 commit into
Conversation
… 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.
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Switching the
efficientnet-lite0backbone to timm (tf_efficientnet_lite0) silently dropped three model-level defaults that the previousefficientnet_lite_pytorchpackage applied. Together they regressed FixMatch fine-tuning quality substantially. This PR restores all three:1/sqrt(fan_in+fan_out)scale, tuned for the 1000-class ImageNet head — roughlynum_classestoo 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 FixMatchp_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 aftertimm.create_model(...).cfg.bn_momentumnever applied. It is computed (= 1 - ema_m, ~0.01) and validated but never set on the model, so BatchNorm ran at timm's0.1default (~10x too fast for our small batch size). Fix: apply it to bothtrain_modelandeval_modelBatchNorm layers.cfg.seednever applied. The seed was defined and validated butset_seedswas 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)
Test plan
ruff check+ruff formatclean on changed filestests/unit/test_net_builder.py::test_efficientnet_classifier_head_init_is_pytorch_scalepassestests/unit/test_fixmatch.pyandtests/unit/test_import.pystill passanomaly_match.pipeline.sessionimports cleanly