DINOv2 Linear Classifier

Train a linear classifier on top of a frozen DINOv2 backbone, and optionally adapt the DINOv2 features to your dataset first.

• dino_training.py — train the classifier on the frozen DINOv2 features.
• dino_adaptation.py — adapt the features semi-supervised, then train the classifier on them.
• dino_inference.py — run inference with either checkpoint.

Data

One folder per class:

data/
  classA/img000.png, img001.png, ...
  classB/img000.png, img001.png, ...

dino_adaptation.py optionally takes a flat folder of unlabeled images, and dino_inference.py a flat folder of images to predict.

Setup

Using pip:

python -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Using conda:

conda env create -f environment.yml
conda activate dinov2_classification

Training

python dino_training.py \
    --data_dir /path/to/data \
    --train_percentage 0.8 \
    --save_dir ./outputs \
    --num_epochs 30 \
    --batch_size 64 \
    --lr 1e-4 \
    --device cuda:0 \
    --model_name dinov2_vits14_reg \
    --early_stopping_patience 10 \
    --seed 42

Writes outputs/best_model.pth, the training curves and a validation report.

Adaptation

The backbone stays frozen and a small projector is trained on top of it with three losses:

1. self-supervised (NT-Xent): two augmented views of the same image attract, every other image repels,
2. supervised contrastive: all views of the same class attract, and they update one prototype per class,
3. prototype pseudo-label: unlabeled images close enough to a prototype are pulled towards it.

The linear classifier is then trained on the adapted features.

python dino_adaptation.py \
    --data_dir /path/to/data \
    --unlabeled_dir /path/to/unlabeled/images \
    --train_percentage 0.8 \
    --save_dir ./outputs_adapted \
    --adapt_epochs 15 \
    --num_epochs 30 \
    --batch_size 64 \
    --lambda_sup 1.0 \
    --lambda_proto 0.5 \
    --device cuda:0 \
    --model_name dinov2_vits14_reg \
    --seed 42

Writes outputs_adapted/adapted_model.pth (projector + classifier), the adaptation curves, the training curves and a validation report.

Inference

python dino_inference.py \
    --inference_dir /path/to/inference/images \
    --checkpoint_path ./outputs/best_model.pth \
    --device cuda:0 \
    --out_file predictions.json

The same command works with ./outputs_adapted/adapted_model.pth: when the checkpoint holds a projector, it is used automatically.