Train a Classifier

trainer.train is the loop used for SimpleCNN, EfficientNet-B0 and the DINOv2 head: AdamW on trainable parameters, cosine schedule down to 1 % of the initial learning rate, cross-entropy with label smoothing 0.1, validation balanced accuracy after each epoch, best checkpoint saved to checkpoints/<model_name>_best.pt and reloaded at the end.

Step 1 — Data

from src.astrovision import load_galaxy10_splits
train_loader, val_loader, test_loader = load_galaxy10_splits(batch_size=64, num_workers=0)

The split is a seed-42 random permutation (Data Preparation).

Step 2 — SimpleCNN from scratch

from src.astrovision import build_model, train
model = build_model('simplecnn')
model = train(model, train_loader, val_loader,
              model_name='simplecnn_galaxy10', epochs=30, lr=1e-3, use_wandb=False, device='cuda')

Reference recipe: 30 epochs at 1e-3 (see Baseline CNN).

Step 3 — EfficientNet-B0 in two phases

effnet = build_model('efficientnet', freeze_backbone=True)       # phase 1: head only
effnet = train(effnet, train_loader, val_loader,
               model_name='efficientnet_galaxy10_warmup', epochs=5, lr=1e-3, use_wandb=False)

for p in effnet.parameters():                                    # phase 2: unfreeze everything
    p.requires_grad = True
effnet = train(effnet, train_loader, val_loader,
               model_name='efficientnet_galaxy10_finetune', epochs=25, lr=1e-4, use_wandb=False)

The reference recipe uses 5 warm-up epochs and 25 fine-tuning epochs (EfficientNet-B0).

Step 4 — Evaluate on the test set

trainer._run_epoch returns loss, acc and balanced_acc; reuse it for the test loader:

import torch.nn as nn
from src.astrovision.trainer import _run_epoch
m = _run_epoch(model, test_loader, None, nn.CrossEntropyLoss(), 'cuda', train=False)
print(m['acc'], m['balanced_acc'])

DINOv2 and fusion

The DINOv2 classifier (DINOv2Classifier) is defined in notebook 04 and the fusion model has its own loop (fusion.train_fusion); neither goes through build_model. Follow the notebooks: fine-tuned DINOv2, late fusion.

Troubleshooting

Symptom	Likely cause
`OSError` from `DataLoader` on Windows	use `num_workers=0`
CUDA out of memory with DINOv2	the notebook uses batch size 32 and FP16 autocast
W&B login prompt	pass `use_wandb=False`
Rare classes never predicted	expected for SimpleCNN (Baseline CNN)