libauc.trainer

The trainer API provides a high-level training interface for classification tasks, including YAML-driven training entry points, callback-based loops, and dataset/model/loss/optimizer orchestration.

Training Pipeline

Module and Class Overview

Name	Description
libauc.trainer.helpers	Utility functions for building evaluation metrics
libauc.trainer.run_image_trainer	CLI entry point for image / tabular classification tasks
libauc.trainer.run_graph_trainer	CLI entry point for graph neural network (GNN) tasks
libauc.trainer.core.callbacks	Callback system for hooking into training lifecycle events
libauc.trainer.core.image_trainer	Core training loop for image and tabular classification models
libauc.trainer.core.graph_trainer	Training loop for graph neural networks, extending Trainer
libauc.trainer.config.args	Dataclass holding all resolved training hyperparameters
libauc.trainer.config.spaces	YAML config parser and default-value resolver
libauc.trainer.data.datasets	Dataset loaders and the IndexedDataset wrapper

libauc.trainer.helpers

libauc.trainer.run_image_trainer

Entry point for image / tabular classification training. It reads a YAML config file, merges it with built-in defaults, builds all components (dataset, metric function, TrainingArguments), and hands off to Trainer.

python -m libauc.trainer.run_image_trainer --config_file config.yaml

Any TrainingArguments field can be overridden directly on the command line, e.g. --epochs 50 --batch_size 64.

libauc.trainer.run_graph_trainer

Entry point for graph neural network (GNN) training. Mirrors run_image_trainer but uses GraphTrainer and a GNN-specific model config (emb_dim, num_layers, etc.).

python -m libauc.trainer.run_graph_trainer --config_file config.yaml

libauc.trainer.core.callbacks

class TrainerState

State object to track training progress.

class CLICallback

Console and Weights & Biases logging callback.

On on_train_begin it initialises a W&B run (silently falls back to console-only when W&B is not installed) and pretty-prints the full TrainingArguments config.

On on_epoch_end it:

• appends a structured entry to state.train_log;

• renders a progress bar (verbose=1) or a per-epoch line (verbose=2) to stdout;

• ships the flat log dict to W&B via wandb.log.

On on_train_end it prints a training summary (best validation and test scores) and calls wandb.finish().

Note

This callback takes no constructor arguments. All configuration is read from TrainingArguments at runtime via the args parameter passed to each lifecycle hook.

Note

W&B logging is silently disabled when wandb is not installed or when wandb.log raises an exception.

on_epoch_begin(args: TrainingArguments, state: TrainerState, **kwargs)

Apply learning-rate / regulariser decay if the epoch is scheduled.

Checks whether state.epoch is listed in args.decay_epochs. If so, decays by a fixed factor of 10:

on_epoch_end(args: TrainingArguments, state: TrainerState, **kwargs)

Record metrics, update the console display, and log to W&B.

1. Appends a structured record {metrics, epoch, lr, train_loss} to state.train_log for later retrieval (e.g. by on_train_end).

2. Console output (controlled by args.verbose):

   • verbose=1 — overwrites a single progress-bar line in place using \r, showing a block-character bar, loss, all eval metrics, and the current LR. A newline is printed only on the final epoch so the bar persists after training.

   • verbose=2 — prints one pipe-separated line per epoch (Epoch N/T | Loss: X | AUROC: Y | LR: Z).

   • verbose=0 — no console output.

3. W&B — ships the flat log dict to wandb.log (no-op if W&B is unavailable).

4. Increments state.epoch.

on_step_end(args: TrainingArguments, state: TrainerState, **kwargs)

Increment the global step counter.

Increments state.step by 1 after every optimizer update, keeping a running total of training steps across all epochs.

on_train_begin(args: TrainingArguments, state: TrainerState, **kwargs)

Initialise W&B and print the training config.

Attempts to start a W&B run using args.project_name and args.experiment_name with the full TrainingArguments dict as the run config. If wandb is not installed the run is silently skipped and self._use_wandb is set to False so all subsequent W&B calls are no-ops.

When args.verbose != 0, pretty-prints the resolved config between two ==== separator lines so the user can confirm all hyperparameters before training starts.

on_train_end(args: TrainingArguments, state: TrainerState, **kwargs)

Print a training summary and close the W&B run.

1. Prints a "Training complete." separator when verbose != 0.

2. Calls wandb.finish() to flush and close the W&B run.

3. Scans state.train_log to find the best epoch (highest value of the first metric on the first eval split) and reports:

   • 1 eval split — best validation score.

   • 2+ eval splits — best validation score + average of all remaining test split scores at that epoch.

4. Stores the summary dict under state.train_summary.

class TrainerCallback

Base class for training lifecycle callbacks.

Every method is a no-op by default, so subclasses only need to override the hooks they care about. Instances are registered with CallbackHandler, which calls each hook in registration order and forwards a consistent set of keyword arguments (model, optimizer, loss_fn, plus any extra kwargs the Trainer supplies for that event).

Lifecycle order during a typical training run:

on_init_end
on_train_begin
  for each epoch:
    on_epoch_begin
      for each step:
        on_step_begin
        on_step_end
    on_evaluate
    on_epoch_end
  [on_save — called periodically inside the epoch loop]
on_train_end

All callback methods are optional and can be overridden in subclasses.

on_epoch_begin(args: TrainingArguments, state: TrainerState, **kwargs)

Called at the start of each epoch. No-op in base class.

on_epoch_end(args: TrainingArguments, state: TrainerState, **kwargs)

Called at the end of each epoch. No-op in base class.

on_evaluate(args: TrainingArguments, state: TrainerState, **kwargs)

Called after each evaluation pass. No-op in base class.

on_init_end(args: TrainingArguments, state: TrainerState, **kwargs)

Called at the end of trainer initialization. No-op in base class.

on_log(args: TrainingArguments, state: TrainerState, **kwargs)

Called after metrics are logged. No-op in base class.

on_predict(args: TrainingArguments, state: TrainerState, metrics, **kwargs)

Called after a prediction pass. No-op in base class.

on_prediction_step(args: TrainingArguments, state: TrainerState, **kwargs)

Called after each prediction step. No-op in base class.

on_save(args: TrainingArguments, state: TrainerState, **kwargs)

Called after a checkpoint is saved. No-op in base class.

on_step_begin(args: TrainingArguments, state: TrainerState, **kwargs)

Called before each optimizer step. No-op in base class.

on_step_end(args: TrainingArguments, state: TrainerState, **kwargs)

Called after each optimizer step. No-op in base class.

on_substep_end(args: TrainingArguments, state: TrainerState, **kwargs)

Called after each gradient-accumulation sub-step. No-op in base class.

on_train_begin(args: TrainingArguments, state: TrainerState, **kwargs)

Called once before the first epoch. No-op in base class.

on_train_end(args: TrainingArguments, state: TrainerState, **kwargs)

Called once after the last epoch. No-op in base class.

class CallbackHandler(callbacks: List[TrainerCallback], model, optimizer, loss_fn)

Multiplexer that owns a list of TrainerCallback instances and fans out every lifecycle event to each of them in registration order.

CallbackHandler itself inherits from TrainerCallback so it can be used polymorphically, but its primary role is orchestration rather than providing hook implementations of its own.

Parameters:

• callbacks (list[TrainerCallback]) – Initial callback list.

• model – The model being trained.

• optimizer – The active optimizer.

• loss_fn – The active loss function.

add_callback(callback)

Add a callback to the handler.

property callback_list

Get a string representation of all callbacks.

pop_callback(callback)

Remove and return a callback.

remove_callback(callback)

Remove a callback without returning it.

libauc.trainer.core.image_trainer

class Trainer(train_args: TrainingArguments, model_cfg: dict, train_dataset: Dataset, eval_dataset: List[Dataset] | None = None, metric: Callable[[Tensor, Tensor], Mapping[str, float]] | None = None, callbacks: List[TrainerCallback] | None = None)

Full training loop for image-classification models supported by libauc.

Trainer wires together a model, an AUC-specific loss function, a libauc optimizer, dual/tri-sampled data loaders, and an optional evaluation pipeline behind a unified train() entry point. Progress is surfaced through a CallbackHandler so any number of TrainerCallback subclasses can observe or alter the training loop without touching Trainer internals.

The class is intentionally thin: heavy lifting (data loading, model construction, loss/optimizer instantiation) is delegated to private helpers so subclasses like GraphTrainer can override only the parts they need.

Parameters:

• train_args (TrainingArguments) – Fully populated training configuration produced by TrainingArguments.

• model_cfg (dict) – Architecture config forwarded to _build_model(). Must contain at least a "name" key. Built-in architectures: resnet20, resnet18, densenet121. Any HuggingFace model repo ID containing a / is also accepted (e.g. "google/vit-base-patch16-224", "openai/clip-vit-base-patch32"). For HF models the AutoImageProcessor is applied inside the trainer — the dataset needs no HF-specific transforms.

• train_dataset (Dataset) – PyTorch Dataset for the training split. Must expose a .targets attribute (list or array of labels).

• eval_dataset (list[Dataset], optional) – One or more evaluation datasets. None disables evaluation (default: None).

• metric (callable, optional) – (y_true, y_pred) -> dict[str, float] function returned by build_metric(). None disables metric computation (default: None).

• callbacks (list[TrainerCallback], optional) – Callbacks invoked at every lifecycle hook. When None the handler is created with an empty list (default: None).

Example:

>>> from trainer.config.args import TrainingArguments
>>> from trainer.core.image_trainer import Trainer
>>> from trainer.core.callbacks import CLICallback
>>> train_args = TrainingArguments(
...     optimizer="PESG", optimizer_kwargs={"lr": 0.1},
...     loss="AUCMLoss", loss_kwargs={"margin": 1.0},
...     SEED=42, batch_size=128, eval_batch_size=128,
...     sampling_rate=0.5, epochs=50, decay_epochs=[],
...     num_workers=2, output_path="./output", num_tasks=1,
...     resume_from_checkpoint=False, save_checkpoint_every=5,
...     project_name="libauc", experiment_name="demo", verbose=1,
... )
>>> trainer = Trainer(
...     train_args=train_args,
...     model_cfg={"name": "resnet18"},
...     train_dataset=train_ds,
...     eval_dataset=[val_ds],
...     metric=metric_fn,
...     callbacks=[CLICallback()],
... )
>>> log = trainer.train()

HuggingFace model — no dataset changes needed, the trainer preprocesses internally:

>>> trainer = Trainer(
...     train_args=train_args,
...     model_cfg={
...         "name": "google/vit-base-patch16-224",
...     },
...     train_dataset=train_ds,
...     eval_dataset=[val_ds],
...     metric=metric_fn,
... )
>>> log = trainer.train()

add_callback(callback)

Add a callback to the trainer.

evaluate(loader, model)

Evaluate model on a given data loader.

Parameters:

• loader – Data loader for evaluation

• model – Model to evaluate

Returns:

Tuple of (dictionary of evaluation metrics, test_true, test_pred)

evaluate_loop(model)

Evaluate model on all evaluation datasets.

Parameters:

model – Model to evaluate

Returns:

Tuple of (dictionary of metrics from all evaluation datasets, test_true, test_pred) test_true and test_pred are from the first evaluation dataset, or None if no eval datasets

get_latest_checkpoint(output_path: str)

load_checkpoint(checkpoint_path: str)

save_checkpoint(checkpoint_path: str)

train()

Main training loop.

Returns:

List of training logs with metrics for each epoch

libauc.trainer.core.graph_trainer

class GraphTrainer(train_args: TrainingArguments, model_cfg: dict, train_dataset, eval_dataset: List | None = None, metric: Callable[[...], Mapping[str, float]] | None = None, callbacks: List[TrainerCallback] | None = None, decay_epochs: List[int] | None = None, decay_factor: float = 10.0, train_eval_dataset=None)

Training loop for graph neural networks built with libauc’s GNN model zoo.

GraphTrainer extends Trainer with graph-aware overrides:

• _build_model() — looks up the requested GNN architecture in _GNN_REGISTRY and constructs it via libauc.models, then infers whether the model expects edge features (supports_edge_attr).

• _get_train_dataloader() / _get_eval_dataloader() — use torch_geometric.loader.DataLoader instead of the standard PyTorch one, while keeping the same DualSampler for positive/negative balancing.

• _forward() — dispatches to the correct GNN forward signature (with or without edge_attr).

• train() — adds optional learning-rate decay at specified epochs via optimizer.update_lr.

Supported GNN architectures

gcn, gin, gine, graphsage, gat, mpnn, deepergcn, pna

param train_args:

Training configuration.

type train_args:

TrainingArguments

param model_cfg:

GNN model configuration. Required key: name (one of the architectures listed above). Optional keys: emb_dim (default 256), num_layers (default 5), graph_pooling, dropout, atom_features_dims, bond_features_dims, act, norm, jk, v2 (GAT-only), aggr / t / learn_t / p / learn_p / block (DeeperGCN-only), pretrained (bool), pretrained_path (str).

type model_cfg:

dict

param train_dataset:

PyG-compatible graph dataset (train split).

param eval_dataset:

PyG-compatible graph datasets for evaluation splits (default: None).

type eval_dataset:

list, optional

param metric:

(y_true, y_pred) -> dict[str, float]

type metric:

callable, optional

param callbacks:

Training callbacks.

type callbacks:

list[TrainerCallback], optional

param decay_epochs:

Epoch indices at which optimizer.update_lr(decay_factor=decay_factor) is called (default: no decay).

type decay_epochs:

list[int], optional

param decay_factor:

LR divisor at each decay epoch (default: 10.0).

type decay_factor:

float

param train_eval_dataset:

Optional dataset for an unbiased train-split evaluation; falls back to train_dataset when None.

Example:

>>> trainer = GraphTrainer(
...     train_args=train_args,
...     model_cfg={"name": "gin", "emb_dim": 300},
...     train_dataset=train_ds,
...     eval_dataset=[val_ds, test_ds],
...     metric=metric_fn,
...     callbacks=[CLICallback()],
...     decay_epochs=[100, 150],
...     decay_factor=10.0,
... )
>>> log = trainer.train()

evaluate(loader, model)

Override base Trainer.evaluate() to use the GNN forward pass.

Parameters:

• loader (PyG DataLoader)

• model (GNN model)

Return type:

(metrics_dict, y_true, y_pred)

train()

GNN training loop.

Steps each epoch:

1. Optional LR decay if epoch is in decay_epochs.

2. Forward / backward over the training loader.

3. Evaluation on the training split (unbiased loader).

4. Evaluation on all registered eval loaders.

5. Callbacks and periodic checkpointing.

Returns:

list

Return type:

training log produced by the state / callback system

libauc.trainer.config.args

class TrainingArguments(**kwargs)

Container for all hyperparameters and settings that govern a single training run.

All fields map one-to-one to keys in the training section of a YAML config file and can be overridden from the CLI via apply_cli_overrides.

Parameters:

• optimizer (str) – Name of the libauc optimizer class, e.g. "PESG", "PDSCA", "SOAP".

• optimizer_kwargs (dict) – Extra keyword arguments forwarded verbatim to the optimizer constructor (e.g. lr, momentum, weight_decay).

• loss (str) – Name of the loss-function class, e.g. "AUCMLoss", "CompositionalAUCLoss". Looked up first in libauc.losses, then torch.nn.

• loss_kwargs (dict) – Extra keyword arguments forwarded verbatim to the loss constructor.

• SEED (int) – Global random seed for NumPy, PyTorch and cuDNN (default: 42).

• batch_size (int) – Mini-batch size for training (default: 128).

• eval_batch_size (int) – Mini-batch size for evaluation (default: 128).

• sampling_rate (float) – Positive-class sampling rate passed to DualSampler / TriSampler (default: 0.5).

• epochs (int) – Total number of training epochs (default: 50).

• decay_epochs (list) – Epoch indices (or fractional multiples of epochs) at which the learning-rate / regulariser is decayed. Floats are converted to int(f * epochs) at construction time.

• num_workers (int) – Number of DataLoader worker processes (default: 2).

• output_path (str) – Root directory for checkpoints and logs (default: "./output").

• num_tasks (int) – Number of output tasks / classes. 1 → binary; ≥ 3 → multi-label with TriSampler.

• resume_from_checkpoint (bool) – Whether to resume from the latest checkpoint found in output_path/experiment_name (default: True).

• save_checkpoint_every (int) – Save a checkpoint every N epochs (default: 5).

• project_name (str) – Weights & Biases project name (default: "libauc").

• experiment_name (str) – Weights & Biases run name; also used as the checkpoint sub-directory.

• verbose (int) – Verbosity level. 0 = silent; 1 = progress bar; 2 = one line per epoch (default: 1).

Example:

>>> args = TrainingArguments(
...     optimizer="PESG",
...     optimizer_kwargs={"lr": 0.1, "momentum": 0.9},
...     loss="AUCMLoss",
...     loss_kwargs={"margin": 1.0},
...     SEED=42,
...     batch_size=128,
...     eval_batch_size=128,
...     sampling_rate=0.5,
...     epochs=50,
...     decay_epochs=[],
...     num_workers=2,
...     output_path="./output",
...     num_tasks=1,
...     resume_from_checkpoint=True,
...     save_checkpoint_every=5,
...     project_name="libauc",
...     experiment_name="my_experiment",
...     verbose=1,
... )

parse_defaultconfig(type_name: str, multilabel: bool = False, kwargs: dict = {})

Resolve a loss or optimizer name to its canonical {optimizer, loss} configuration dict by looking up the corresponding spaces class.

The mapping covers every loss/optimizer pair supported by libauc:

type_name	Space class
AUCMLoss / PESG	AUCMLossSpace (MultiLabelAUCMLossSpace when multilabel=True)
CompositionalAUCLoss / PDSCA	CompositionalAUCLossSpace
APLoss / SOAP	APLossSpace (mAPLossSpace when multilabel=True)
pAUC_CVaR_Loss / SOPA / pAUCLoss mode SOPA	pAUC_CVaR_LossSpace (MultiLabel… variant)
pAUC_DRO_Loss / SOPAs / pAUCLoss mode 1w	pAUC_DRO_LossSpace (MultiLabel… variant)
tpAUC_KL_Loss / SOTAs / pAUCLoss mode 2w	tpAUC_KL_LossSpace (MultiLabel… variant)
tpAUC_CVaR_loss / STACO	tpAUC_CVaR_lossSpace
NDCGLoss / SONG	NDCGLossSpace
CrossEntropyLoss / SGD	SGDSpace
Adam	AdamSpace
BCELoss	BCELossSpace

Parameters:

• type_name (str) – Name of the loss or optimizer class.

• multilabel (bool) – When True, selects the multi-label variant of the space if one exists (default: False).

• kwargs (dict) – Additional keyword arguments for the loss/optimizer, used to disambiguate pAUCLoss by its mode key.

Returns:

{"optimizer": <optimizer_cfg>, "loss": <loss_cfg>} where each config is a dict with at least a "type" key and a "space" key containing the hyperparameter search space.

Return type:

dict

Raises:

ValueError – If type_name is not recognised.

Example:

>>> cfg = parse_defaultconfig("AUCMLoss", multilabel=False)
>>> cfg["optimizer"]["type"]
'PESG'
>>> cfg["loss"]["type"]
'AUCMLoss'

libauc.trainer.config.spaces

class APLossSpace

loss = {'space': {'gamma': {'default': 0.9, 'val': (0.0, 1.0)}, 'margin': {'default': 1.0, 'val': [0.6, 0.8, 1.0]}}, 'type': 'APLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'SOAP'}

class AUCMLossSpace

loss = {'space': {'margin': {'default': 1.0, 'val': [0.6, 0.8, 1.0]}}, 'type': 'AUCMLoss'}

optimizer = {'space': {'epoch_decay': {'default': 0.002, 'val': (0.0, 0.01)}, 'lr': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'PESG'}

class AdamSpace

loss = {'space': {}, 'type': 'CrossEntropyLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'Adam'}

class BCELossSpace

loss = {'space': {}, 'type': 'BCELoss'}

optimizer = {}

class CompositionalAUCLossSpace

loss = {'space': {'k': {'default': 1, 'val': [1, 2, 4]}, 'margin': {'default': 1.0, 'val': [0.6, 0.8, 1.0]}}, 'type': 'CompositionalAUCLoss'}

optimizer = {'space': {'epoch_decay': {'default': 0.002, 'val': (0.0, 0.01)}, 'lr': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'PDSCA'}

class MultiLabelAUCMLossSpace

loss = {'space': {'margin': {'default': 1.0, 'val': [0.6, 0.8, 1.0]}}, 'type': 'MultiLabelAUCMLoss'}

optimizer = {'space': {'epoch_decay': {'default': 0.002, 'val': (0.0, 0.01)}, 'lr': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'PESG'}

class MultiLabelpAUC_CVaR_LossSpace

loss = {'space': {'beta': {'val': 0.2}, 'eta': {'default': 0.1, 'log': True, 'val': (0.01, 10)}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOPA'}}, 'type': 'MultiLabelpAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SOPA'}

class MultiLabelpAUC_DRO_LossSpace

loss = {'space': {'Lambda': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}, 'gamma': {'default': 0.9, 'val': (0.0, 1.0)}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOPAs'}}, 'type': 'MultiLabelpAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'SOPAs'}

class MultiLabeltpAUC_KL_LossSpace

loss = {'space': {'Lambda': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}, 'gammas': {'default': (0.9, 0.9), 'val': [(0.1, 0.1), (0.5, 0.5), (0.9, 0.9)]}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOTAs'}, 'tau': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}}, 'type': 'MultiLabelpAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SOTAs'}

class NDCGLossSpace

loss = {'space': {'eta0': {'default': 0.01, 'log': True, 'val': (0.001, 0.1)}, 'gamma0': {'default': 0.9, 'val': (0.0, 1.0)}, 'gamma1': {'val': 0.9}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'sigmoid_alpha': {'default': 2.0, 'val': (1.0, 2.0)}}, 'type': 'NDCGLoss'}

optimizer = {'space': {'lr': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SONG'}

class SGDSpace

loss = {'space': {}, 'type': 'CrossEntropyLoss'}

optimizer = {'space': {'lr': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0, 'val': [0, 0.9]}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SGD'}

class mAPLossSpace

loss = {'space': {'gamma': {'default': 0.9, 'val': (0.0, 1.0)}, 'margin': {'default': 1.0, 'val': [0.6, 0.8, 1.0]}}, 'type': 'mAPLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'SOAP'}

class pAUC_CVaR_LossSpace

loss = {'space': {'beta': {'val': 0.2}, 'eta': {'default': 0.1, 'log': True, 'val': (0.01, 10)}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOPA'}}, 'type': 'pAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SOPA'}

class pAUC_DRO_LossSpace

loss = {'space': {'Lambda': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}, 'gamma': {'default': 0.9, 'val': (0.0, 1.0)}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOPAs'}}, 'type': 'pAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 1e-05, 'val': (0.0, 0.0002)}}, 'type': 'SOPAs'}

class tpAUC_CVaR_lossSpace

loss = {'space': {'alpha': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'beta_0': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'beta_1': {'default': 0.1, 'log': True, 'val': (0.0001, 0.1)}, 'theta_0': {'default': 0.5, 'val': [0.3, 0.5, 0.7]}, 'theta_1': {'default': 0.5, 'val': [0.3, 0.5, 0.7]}, 'threshold': {'default': 0.5, 'val': [0.3, 0.5, 0.7]}}, 'type': 'tpAUC_CVaR_loss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.01)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'STACO'}

class tpAUC_KL_LossSpace

loss = {'space': {'Lambda': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}, 'gammas': {'default': (0.9, 0.9), 'val': [(0.1, 0.1), (0.5, 0.5), (0.9, 0.9)]}, 'margin': {'default': 1.0, 'val': [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]}, 'mode': {'val': 'SOTAs'}, 'tau': {'default': 1.0, 'log': True, 'val': (0.1, 10.0)}}, 'type': 'pAUCLoss'}

optimizer = {'space': {'lr': {'default': 0.001, 'log': True, 'val': (0.0001, 0.1)}, 'momentum': {'default': 0.9, 'val': (0.8, 0.99)}, 'weight_decay': {'default': 0, 'val': (0.0, 0.0002)}}, 'type': 'SOTAs'}

libauc.trainer.data.datasets

class ChemicalDataset(dataset, class_id)

OGB molecular graph dataset filtered to a single task column.

class GraphDataset(name, root='dataset', transform=None, pre_transform=None, meta_dict=None)

PygGraphPropPredDataset with integer-index support.

class ImageDataset(images, targets, image_size=32, crop_size=30, mode='train')

In-memory image dataset with train/test augmentation presets.

class IndexedDataset(dataset, class_id=None)

Wraps an existing dataset to return (image, target, index) tuples.

Optionally selects a single column from multi-label targets when class_id is given.

class MedicalImageCSVDataset(source, image_root: str, image_col: str, label_col: str, transform)

General-purpose CSV-backed medical image dataset.

Expects a CSV (or DataFrame) with at least an image path column and a binary label column. Image paths may be relative (resolved against image_root) or absolute.

Parameters:

• source – Path to a metadata CSV or a pandas.DataFrame with the required columns already loaded. Passing a DataFrame avoids writing temporary files.

• image_root – Directory that relative image paths are resolved against. Ignored for absolute paths.

• image_col – Column name containing the image filename / path.

• label_col – Column name containing the binary label (0 / 1).

• transform – torchvision transform applied to each PIL image.

load_dataset(name: str, splits: list, **kwargs)

Load a dataset by name and return train + eval splits.

Parameters:

• name – Dataset identifier (case-insensitive).

• splits – Evaluation splits to return, e.g. ["val", "test"].

• **kwargs – Extra dataset-specific keyword arguments from the config.

Returns:

(train_dataset, eval_datasets) — both are torch.utils.data.Dataset instances whose __getitem__ yields (data, label, index) tuples, as expected by the Trainer.