Vit Flower Classifier

ViT-B/16 Flower Classifier

Fine-tuned torchvision.models.vit_b_16 (ImageNet-1K pretrained) for 102-class flower classification on the Oxford-102 Flowers dataset, with the full backbone unfrozen during fine-tuning. Achieves 0.9837 accuracy / 0.9722 macro-F1 on a held-out test split — the best of every architecture evaluated for this project.

Note: this card describes a new checkpoint, retrained after a data leak in the training pipeline's split logic was fixed. The weights previously published under this repo were leak-era and were selected on a corrupted validation signal — see Training history. If you pinned an earlier revision, re-download.

Recommended when macro-F1 is the priority and the size budget is not binding — offline batch labeling, research baselines, or any deployment where 343.5 MB of weights is acceptable. For most serving, ConvNeXt-Tiny is the better trade: it scores 0.9694 macro-F1 against this model's 0.9722, at 111.6 MB instead of 343.5 MB.

Usage

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torchvision import models

weights_path = hf_hub_download(repo_id="bengid/vit-flower-classifier", filename="vit-flower-classifier.safetensors")

model = models.vit_b_16(weights=None)
model.heads[-1] = torch.nn.Linear(model.heads[-1].in_features, 102)
model.load_state_dict(load_file(weights_path, device="cpu"))
model.eval()

# preprocessing: Resize((256, 256)) -> CenterCrop(224) -> ToTensor -> Normalize(dataset mean/std)
# note: the resize is to a fixed 256x256 square, NOT shortest-side-256 -- aspect ratio is not preserved
# see src/utils.py:get_transforms() in the training repo for the exact pipeline
# (https://github.com/ben-gid/flowers/blob/main/src/utils.py)

Training Data

Oxford-102 Flowers — 8,189 images across 102 flower species, downloaded via torchvision.datasets.Flowers102. Class-weighted CrossEntropyLoss was used to correct for the dataset's uneven per-class image counts.

Training Procedure

Single-stage fine-tune with a two-phase backbone unfreeze callback (BackboneFinetuning): the ViT backbone starts frozen (only the classification head trains), then unfreezes at a fixed epoch with its own, lower learning rate and a separate parameter group — unlike this project's original (v1) EfficientNet-B0 model, which only ever unfroze its last 3 backbone blocks.

Trained under the corrected-split-retrain MLflow experiment, i.e. after the split bug described in Training history was fixed. Checkpoint selection and early stopping were therefore driven by a genuine held-out validation signal.

Hyperparameters

Parameter	Value
Optimizer	AdamW
LR scheduler	Cosine annealing (T_max=50, eta_min=1e-06)
Head LR (before unfreeze)	0.001
Head LR (after unfreeze)	0.0003
Backbone LR (after unfreeze)	2e-05
Unfreeze epoch	5
Max epochs	50
Batch size	64
Effective batch size	256
Gradient accumulation	4
Precision	bf16-mixed
Weight decay	0.01
Early stopping patience	8

Evaluation

Oxford-102 split 70/15/15 by random_split(seed=42) → 5,733 train / 1,228 val / 1,228 test. Val and test are disjoint from train and from each other.

Metric	Validation	Test
Accuracy	0.9837	0.9837
Macro F1	0.9737	0.9722
Loss	0.0797	0.0762

Property	Value
Parameters	85,877,094
Model size	343.5 MB
Checkpoint size	1030.7 MB
Mean latency	5.39 ms
p95 latency	5.68 ms

Latency measured on an NVIDIA RTX 5070 (cuda:0) at batch size 1, with torch.cuda.synchronize() after every forward pass. These figures do not transfer to CPU. An earlier CPU benchmark of the same architectures put ViT-B/16 at roughly 3x the latency of the other two, while ConvNeXt-Tiny and EfficientNetV2-S landed within a few milliseconds of each other — close enough that two runs of that benchmark disagreed on which was faster. Benchmark on your own serving hardware before treating any of this as a ranking.

Macro F1 sits below accuracy on both splits, which is the signature of uneven per-class performance on a dataset whose class counts range from 40 to 258 images — the rare classes are where the misses are. This model has the narrowest accuracy-to-F1 spread of the three, though it leads ConvNeXt-Tiny only marginally on that measure. Per-class F1 is computed but not currently exported; see test_per_class_f1 in src/classifier.py.

Training history

An earlier version of this card claimed 1.0 accuracy / 1.0 F1. That was a measurement bug, not a real result. FlowerDataModule.setup() in src/data.py built all three splits from the same subset:

train_subset, val_subset, test_subset = random_split(full_ds, (0.7, 0.15, 0.15), generator=generator)

self.train_set = SubsetWithTransform(train_subset, self.transform_train)
self.val_set   = SubsetWithTransform(train_subset, self.transform_test)  # bug — should be val_subset
self.test_set  = SubsetWithTransform(train_subset, self.transform_test)  # bug — should be test_subset

val_subset and test_subset were built and then discarded. Validation and test scored the model against the exact images it had trained on — total leakage, not partial. A ~1.0 score was the expected outcome of memorization and carried no information about generalization.

The weights were never contaminated. Training only ever read train_subset, so no held-out image ever reached a gradient. What the bug corrupted was measurement — and, through it, model selection: ModelCheckpoint and EarlyStopping were both driven by val_acc, which was really train accuracy.

This checkpoint replaces the leak-era one. It comes from the corrected-split-retrain experiment, trained after the fix, so every metric above was measured on genuinely held-out data and the checkpoint was selected on a real validation signal.

One finding worth recording: the retrain recovered less than expected. The leak-era ViT checkpoint, re-scored honestly on the same held-out data, reached 0.9637 test macro-F1 against this model's 0.9722. The bug inflated the reported numbers substantially; it cost the weights comparatively little.

Strengths & Weaknesses

Strengths:

• Best accuracy and macro-F1 of every architecture evaluated for this project — attention-based global context handles flowers that differ mainly in overall shape/arrangement rather than local texture.
• Narrowest accuracy-to-macro-F1 spread of the three published models, so its errors are the least concentrated in rare classes — though it is only marginally ahead of ConvNeXt-Tiny here.
• Full-backbone fine-tuning lets every ViT layer adapt to the flower domain, avoiding the ceiling that partial-unfreeze approaches hit.

Weaknesses:

• Largest model in the lineup by a wide margin — 85,877,094 parameters and 343.5 MB of weights, roughly three times ConvNeXt-Tiny, for a macro-F1 margin measured in fractions of a point. For most serving that trade is not worth paying.
• The accuracy lead over ConvNeXt-Tiny is small enough (0.9722 vs 0.9694) that it is within the range a different seed or a longer schedule could move. Treat it as a lead, not a law.
• ViTs are comparatively data-hungry and were historically harder to fine-tune from limited data before full-backbone unfreezing plus a long enough schedule — this model only reaches its ceiling because both were used.
• Not a good fit for edge/mobile deployment given its size.

Limitations

• Closed-set, single-label: trained on exactly 102 Oxford flower species; will confidently misclassify any other flower species, non-flower image, or multi-flower image into one of the 102 known classes — there is no out-of-distribution rejection.
• Fixed input pipeline: expects a 224×224 input produced by resizing to a fixed 256×256 square and center-cropping. The resize does not preserve aspect ratio, so non-square images are distorted before the crop, and off-center subjects can be cropped out of frame.
• No adversarial robustness or calibration guarantees — confidence scores are not calibrated probabilities.
• Reported metrics are on held-out Oxford-102 val/test splits; real-world images (different lighting, backgrounds, camera quality) may perform worse.
• Macro F1 (0.9722) is the number to trust, not accuracy (0.9837) — Oxford-102 is class-imbalanced (40–258 images per class), and the gap between the two means errors concentrate in rare classes. If your use case cares about the long tail, budget for the F1 figure.

Intended Use

Intended uses:

• Flower species identification within the 102 Oxford-102 classes (gardening/botany apps, educational tools, dataset labeling).
• Backend model for this project's v2 /classify API endpoint when macro-F1 is prioritized over model size.

Out-of-scope uses:

• General-purpose plant, object, or scene classification outside the 102 trained species.
• Medical, toxicity, or safety-related plant identification.
• Any use where a wrong classification has safety or financial consequences without human review.

Model Comparison

All figures are test-split metrics on held-out data. The three current models were trained on the corrected splits; the two v1 baselines are listed for historical context.

Model	Test Acc	Test F1	Params	Size (MB)	Best For
SimpleCNN (scratch)	~0.63	-	-	-	historical baseline only
EfficientNet-B0 (v1, partial unfreeze)	>0.93	-	-	-	historical baseline only
ViT-B/16 (this model)	0.9837	0.9722	85,877,094	343.5	maximum macro-F1
ConvNeXt-Tiny	0.9821	0.9694	27,898,566	111.6	default serving choice
EfficientNetV2-S	0.9691	0.9433	20,308,150	81.8	smallest weights

The two v1 models were trained by the older, pre-Lightning pipeline (api/app/v1/flowers/train_scratch.py), which split train/val/test correctly — their numbers were never affected by the leak and are directly comparable to the corrected figures above.

Why the earlier models underperformed

• SimpleCNN (scratch) was trained from randomly initialized weights with no ImageNet pretraining, on a 6-block custom CNN — too little capacity and too little prior visual knowledge to learn 102 fine-grained flower classes from ~8k images alone.
• EfficientNet-B0 (v1) started from ImageNet-pretrained weights but only ever unfroze its last 3 backbone blocks during fine-tuning (see this project's root README.md for the...