Skip to main content

trainer.py

train(model, train_loader, val_loader, model_name='astrovision', epochs=30, lr=1e-3, weight_decay=1e-4, use_wandb=True, device='cuda') → nn.Module​

ArgumentMeaning
modelany nn.Module returning logits
train_loader, val_loaderloaders yielding (images, labels)
model_namecheckpoint stem: weights go to checkpoints/<model_name>_best.pt (directory next to src/, created on import)
epochs, lr, weight_decayAdamW over requires_grad parameters; CosineAnnealingLR(T_max=epochs, eta_min=lr × 0.01)
use_wandblog to W&B project AstroVision (loss, accuracy, balanced accuracy, learning rate per epoch)
devicefalls back to CPU if CUDA is unavailable

Behaviour:

  • loss: CrossEntropyLoss(label_smoothing=0.1);
  • each epoch computes loss, acc and balanced_acc (scikit-learn) on train and validation;
  • the checkpoint is overwritten when validation balanced accuracy improves; at the end the best weights are reloaded into the model, which is returned;
  • the best-epoch criterion is balanced accuracy on the validation set, never the test set.

Not included: mixed precision (the DINOv2 and U-Net notebooks use their own loops with FP16 autocast), early stopping, class re-weighting.

Private helper _run_epoch(model, loader, optimizer, criterion, device, train) returns {'loss', 'acc', 'balanced_acc'}; it is convenient for test-set evaluation (guide).

model = train(model, train_loader, val_loader,
              model_name='efficientnet_galaxy10_finetune', epochs=25, lr=1e-4, use_wandb=False)