Train with LibAUC Trainer
Introduction
The LibAUC Trainer is a high-level training interface that turns a YAML config file into a complete training run — data loading, model construction, loss/optimizer wiring, evaluation, and checkpointing are all handled automatically.
This tutorial walks through a two-stage workflow for AUROC maximization:
Stage |
Goal |
Config |
|---|---|---|
1 |
Warm-up with cross-entropy pretraining |
|
2 |
Fine-tune for AUROC with AUC-aware loss |
|
Note
Pretraining first and then switching to an AUC loss typically yields higher AUROC than training with the AUC loss from scratch, because the model starts from a meaningful feature representation.
Quick Start
# Stage 1 — cross-entropy warm-up
python -m libauc.trainer.run_trainer --config_file ce_config.yaml
# Stage 2 — AUROC optimization
python -m libauc.trainer.run_trainer --config_file aucmloss_config.yaml
Any field in the YAML can be overridden directly on the command line:
python -m libauc.trainer.run_trainer --config_file aucmloss_config.yaml \
--epochs 50 --batch_size 64 --sampling_rate 0.3
Supported Loss / Optimizer Pairings
The loss and optimizer fields must be a compatible pair.
Task |
Loss |
Optimizer |
|---|---|---|
AUROC (binary) |
|
|
AUROC (multi-label) |
|
|
Compositional AUROC |
|
|
Average Precision |
|
|
Partial AUROC (CVaR) |
|
|
Partial AUROC (DRO) |
|
|
Two-way partial AUROC (KL) |
|
|
Two-way partial AUROC (CVaR) |
|
|
NDCG |
|
|
CE pretraining (SGD) |
|
|
CE pretraining (Adam) |
|
|
For parameter details see the libauc.optimizers API reference.
Step 1: CE Pretraining
Create ce_config.yaml:
# ── Dataset: which dataset to load, splits to evaluate, and class-imbalance ratio ──
dataset:
name: cifar10 # see "Supported Datasets" below
eval_splits: [val, test]
kwargs:
imratio: 0.1 # positive-class ratio in the imbalanced training set
# ── Model: architecture, weight initialization, and output format ──────────────────
model:
name: resnet18 # see "Supported Models" below
pretrained: false
num_classes: 1 # 1 → binary; ≥ 3 → multi-label
in_channels: 3
# ── Metrics: evaluation metrics computed on every eval split ───────────────────────
metrics:
- AUROC # AUROC | AUPRC | ACC
# ── Training: hyperparameters, loss, optimizer, and checkpointing ──────────────────
training:
# Experiment metadata
project_name: libauc
experiment_name: resnet18_ce_cifar10
SEED: 2026
# Data loading
epochs: 100
batch_size: 128
eval_batch_size: 256
sampling_rate: 0.5 # positive fraction per batch (DualSampler)
num_workers: 2
decay_epochs: [0.5, 0.75] # fractions of total epochs, or absolute ints
# Loss
loss: BCELoss
loss_kwargs: {}
# Optimizer
optimizer: Adam
optimizer_kwargs:
lr: 1.0e-3
weight_decay: 1.0e-4
# Output and checkpointing
output_path: ./output
resume_from_checkpoint: false
save_checkpoint_every: 5
verbose: 1 # 0 = silent | 1 = progress bar | 2 = one line/epoch
Run it:
python -m libauc.trainer.run_trainer --config_file ce_config.yaml
Expected checkpoint path after training:
./output/resnet18_ce_cifar10/epoch_100.pt
Step 2: AUROC Optimization
Create aucmloss_config.yaml:
# ── Dataset: same split and imbalance ratio as Stage 1 ────────────────────────────
dataset:
name: cifar10
eval_splits: [val, test]
kwargs:
imratio: 0.1
# ── Model: fine-tune from the Stage-1 checkpoint ──────────────────────────────────
model:
name: resnet18
pretrained: true
pretrained_path: "./output/resnet18_ce_cifar10/epoch_100.pt"
num_classes: 1
in_channels: 3
# ── Metrics: evaluation metrics computed on every eval split ───────────────────────
metrics:
- AUROC
# ── Training: AUC-aware loss and PESG optimizer ────────────────────────────────────
training:
# Experiment metadata
project_name: libauc
experiment_name: resnet18_AUCMLoss_cifar10
SEED: 2026
# Data loading
epochs: 100
batch_size: 128
eval_batch_size: 256
sampling_rate: 0.2 # lower ratio is often better for AUC losses
num_workers: 2
decay_epochs: [0.5, 0.75]
# Loss
loss: AUCMLoss
loss_kwargs:
margin: 1.0 # decision boundary margin, typical range [0.6, 1.0]
# Optimizer
optimizer: PESG
optimizer_kwargs:
lr: 0.1
epoch_decay: 0.002
weight_decay: 1.0e-5
momentum: 0.9
# Output and checkpointing
output_path: ./output
resume_from_checkpoint: false
save_checkpoint_every: 5
verbose: 1
Run it:
python -m libauc.trainer.run_trainer --config_file aucmloss_config.yaml
Config Reference
Supported Datasets
The dataset.name key selects the dataset and dataset.kwargs passes
dataset-specific arguments to the loader.
See libauc.trainer.data.datasets.load_dataset
for the full implementation.
|
|
|
|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Supported Models
|
|
Notes |
|---|---|---|
|
3 |
Supports |
|
3 |
Lightweight; designed for CIFAR-scale inputs |
|
3 |
Supports |
When pretrained: true, the Trainer loads a local checkpoint from
pretrained_path and resets the final classification head (fc or
linear) so it can be retrained for the new loss.
When pretrained_remote: true, ImageNet weights are downloaded via libauc’s
model hub (independent of pretrained).
Supported Metrics
|
Behaviour |
|---|---|
|
Full-ROC AUROC. With |
|
Area under the precision-recall curve. |
|
Accuracy at threshold 0.5. |
Example — track both full AUROC and partial AUROC simultaneously:
metrics:
- AUROC
- AUROC
metric_kwargs:
- {} # full AUROC
- {max_fpr: 0.3} # pAUROC at FPR ≤ 0.3
Training Field Reference
Field |
Default |
Description |
|---|---|---|
|
|
Weights & Biases project name |
|
|
Run name; also used as the checkpoint sub-directory |
|
|
Global random seed (NumPy, PyTorch, cuDNN) |
|
|
Total training epochs |
|
|
Mini-batch size during training |
|
|
Mini-batch size during evaluation |
|
|
Positive-class fraction fed to |
|
|
DataLoader worker processes |
|
|
Epochs at which the LR / regulariser is decayed.
Floats are multiplied by |
|
|
Loss class name (see pairings table above) |
|
|
Extra keyword arguments forwarded to the loss constructor |
|
|
Optimizer class name |
|
|
Extra keyword arguments forwarded to the optimizer constructor |
|
|
Root directory for checkpoints |
|
|
Resume from the latest checkpoint in |
|
|
Save a checkpoint every N epochs |
|
|
|
Extending the Trainer
Adding a New Dataset
All dataset logic lives in a single function:
libauc/trainer/data/datasets.py → load_dataset().
Add a new elif branch that returns (train_dataset, eval_datasets).
Both must be torch.utils.data.Dataset subclasses whose __getitem__
yields (data, label, index) tuples.
# libauc/trainer/data/datasets.py
def load_dataset(name: str, splits: List[str], **kwargs) -> Dataset:
...
elif name == "mydata":
root = kwargs.get("root_path", "./data")
# Build train / eval datasets here.
# __getitem__ must return (data, label, index).
train_dataset = MyTrainDataset(root=root)
eval_datasets = []
for split in splits:
if split == "val":
eval_datasets.append(MyValDataset(root=root))
elif split == "test":
eval_datasets.append(MyTestDataset(root=root))
else:
raise NotImplementedError(
f"Split '{split}' not supported for 'mydata'."
)
return train_dataset, eval_datasets
...
Then reference it in your config:
dataset:
name: mydata
eval_splits: [val, test]
kwargs:
root_path: /path/to/mydata
Tip
Use the built-in IndexedDataset wrapper to add the required index
return value to any standard torchvision dataset:
from libauc.trainer.data.datasets import IndexedDataset
import torchvision.datasets as tvd
train_dataset = IndexedDataset(
tvd.ImageFolder(root=os.path.join(root, "train"), transform=train_transform)
)
Adding a New Model
Model construction is handled in
libauc/trainer/core/trainer.py → Trainer._build_model().
Add a new elif branch that instantiates your model and assigns it to
self.model:
# libauc/trainer/core/trainer.py
def _build_model(self, model_cfg: dict):
name = model_cfg.get("name", "").lower()
pretrained = model_cfg.get("pretrained", False)
pretrained_remote = model_cfg.get("pretrained_remote", False)
num_classes = model_cfg.get("num_classes", 1)
in_channels = model_cfg.get("in_channels", 3)
if name == "resnet18":
...
elif name == "mymodel":
from mypackage import MyModel
model = MyModel(num_classes=num_classes, in_channels=in_channels)
else:
raise ValueError(f"Unknown model '{name}'.")
model = model.cuda()
if pretrained:
state_dict = torch.load(model_cfg.get("pretrained_path"), weights_only=False)
if "model_state_dict" in state_dict:
state_dict = state_dict["model_state_dict"]
# Strip the classification head so it is re-initialised
filtered = {k: v for k, v in state_dict.items()
if "fc" not in k and "linear" not in k}
model.load_state_dict(filtered, strict=False)
# Reset the head
if hasattr(model, "fc"):
model.fc.reset_parameters()
self.model = model
Then reference it in your config:
model:
name: mymodel
num_classes: 1
in_channels: 3
Note
The model’s final layer must output raw logits (no sigmoid / softmax).
Pass last_activation=None when using libauc built-in models, as the
loss functions apply their own activation internally.
HuggingFace Integration
This section shows how to use a HuggingFace dataset and a HuggingFace transformer model with the LibAUC Trainer end-to-end, without touching the rest of the codebase.
Install extra dependencies first:
pip install datasets transformers
Step 1 — Register a HuggingFace Dataset
Add a wrapper class and a new branch inside
libauc/trainer/data/datasets.py → load_dataset().
The only contract the Trainer requires is that __getitem__ returns
(image_tensor, float_label, index).
# libauc/trainer/data/datasets.py
from datasets import load_dataset as hf_load_dataset
from torch.utils.data import Dataset
from torchvision import transforms
import numpy as np
class HFImageDataset(Dataset):
"""Wraps a HuggingFace image split for LibAUC Trainer."""
def __init__(self, hf_split, pos_class, transform=None):
self.data = hf_split
self.pos_class = pos_class
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
image = item["img"] # PIL Image
label = np.float32(item["label"] == self.pos_class)
if self.transform:
image = self.transform(image)
return image, label, idx
def load_dataset(name, splits, **kwargs):
...
elif name == "hf_cifar10":
pos_class = kwargs.get("pos_class", 0) # e.g. 0 = airplane
_mean, _std = [0.5] * 3, [0.5] * 3
train_tf = transforms.Compose([
transforms.Resize(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(_mean, _std),
])
eval_tf = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(_mean, _std),
])
raw = hf_load_dataset("cifar10")
train_dataset = HFImageDataset(raw["train"], pos_class, train_tf)
eval_datasets = []
for split in splits:
if split in ("val", "test"):
eval_datasets.append(
HFImageDataset(raw["test"], pos_class, eval_tf)
)
else:
raise NotImplementedError(f"Split '{split}' not supported.")
return train_dataset, eval_datasets
...
Step 2 — Register a HuggingFace Model
Add a new branch inside
libauc/trainer/core/trainer.py → Trainer._build_model().
Define the wrapper class inline so no extra files are needed.
# libauc/trainer/core/trainer.py
def _build_model(self, model_cfg):
name = model_cfg.get("name", "").lower()
num_classes = model_cfg.get("num_classes", 1)
pretrained = model_cfg.get("pretrained", False)
...
elif name == "vit":
from transformers import ViTModel
import torch.nn as nn
class ViTClassifier(nn.Module):
def __init__(self, hf_name, num_classes):
super().__init__()
self.vit = ViTModel.from_pretrained(hf_name)
hidden = self.vit.config.hidden_size
self.head = nn.Linear(hidden, num_classes)
def forward(self, x):
out = self.vit(pixel_values=x)
return self.head(out.pooler_output)
hf_name = model_cfg.get("hf_name", "google/vit-base-patch16-224")
model = ViTClassifier(hf_name, num_classes)
...
model = model.cuda()
if pretrained:
state = torch.load(model_cfg["pretrained_path"], weights_only=False)
state = state.get("model_state_dict", state)
model.load_state_dict(
{k: v for k, v in state.items() if "head" not in k},
strict=False,
)
model.head.reset_parameters()
self.model = model
Step 3 — Write the YAML Configs
Stage 1 — cross-entropy warm-up (ce_hf_config.yaml):
dataset:
name: hf_cifar10
eval_splits: [val, test]
kwargs:
pos_class: 0 # airplane = positive class (binary AUROC)
model:
name: vit
hf_name: google/vit-base-patch16-224
num_classes: 1 # binary → single logit
metrics:
- AUROC
training:
experiment_name: vit_bce_hf_cifar10
epochs: 5
batch_size: 32
eval_batch_size: 64
sampling_rate: 0.5
num_workers: 4
decay_epochs: [0.6, 0.9]
loss: BCELoss
loss_kwargs: {}
optimizer: Adam
optimizer_kwargs:
lr: 2.0e-5 # small LR for fine-tuning
weight_decay: 1.0e-4
output_path: ./output
save_checkpoint_every: 1
verbose: 1
Stage 2 — AUROC optimization (aucm_hf_config.yaml):
dataset:
name: hf_cifar10
eval_splits: [val, test]
kwargs:
pos_class: 0
model:
name: vit
hf_name: google/vit-base-patch16-224
num_classes: 1
pretrained: true
pretrained_path: "./output/vit_bce_hf_cifar10/epoch_5.pt"
metrics:
- AUROC
training:
experiment_name: vit_AUCMLoss_hf_cifar10
epochs: 10
batch_size: 32
eval_batch_size: 64
sampling_rate: 0.2
num_workers: 4
decay_epochs: [0.5, 0.75]
loss: AUCMLoss
loss_kwargs:
margin: 1.0
optimizer: PESG
optimizer_kwargs:
lr: 0.05
epoch_decay: 0.002
weight_decay: 1.0e-5
momentum: 0.9
output_path: ./output
save_checkpoint_every: 1
verbose: 1
Run both stages:
python -m libauc.trainer.run_trainer --config_file ce_hf_config.yaml
python -m libauc.trainer.run_trainer --config_file aucm_hf_config.yaml
Tip
pos_class controls which CIFAR-10 label is treated as the positive class.
Swap in any HuggingFace image dataset by changing hf_load_dataset("cifar10")
and the "img" / "label" field names to match that dataset’s schema.
Use datasets.load_dataset(...).features to inspect the available fields.
Expected Outputs
After both stages finish, your output directory will look like:
./output/
├── resnet18_ce_cifar10/
│ ├── epoch_5.pt
│ ├── ...
│ └── epoch_100.pt ← loaded by aucmloss_config.yaml
└── resnet18_AUCMLoss_cifar10/
├── epoch_5.pt
├── ...
└── epoch_100.pt
Validation and test AUROC scores are printed after every evaluation epoch.