peal.adaptors.group_distributionally_robust_optimization.GroupDRO

class peal.adaptors.group_distributionally_robust_optimization.GroupDRO[source]

Bases: Adaptor

Adaptor that (re)trains a predictor with the GroupDRO objective.

Groups are y * output_size + confounder for PEAL datasets and the dataset’s own group index for the replication datasets. Only SGD (momentum 0.9, weight decay task.criterions["l2"]) and a single ce/bce/ mse/mae objective are supported.

Parameters:

  • adaptor_config (dict or str or Path or AdaptorConfig) – A GroupDROConfig (or yaml path); predictor supplies the architecture, training, data and task configs.

  • model_path (str, optional) – Overrides adaptor_config.model_path; default is the predictor’s model_path with a DRO_ prefix on its last component.

  • model (torch.nn.Module, optional) – Model to adapt; built from the architecture config when None (SequentialModel for replication runs, TorchvisionModel for torchvision_* strings).

  • datasource (str or sequence of Dataset, optional) – Dataset root or (train, val[, test]) datasets; defaults to data_config.dataset_path. Ignored for replication datasets.

  • optimizer – Accepted for interface compatibility with ModelTrainer but unused.

  • criterions – Accepted for interface compatibility with ModelTrainer but unused.

  • logger – Accepted for interface compatibility with ModelTrainer but unused.

  • only_last_layer – Accepted for interface compatibility with ModelTrainer but unused.

  • unit_test_train_loop – Accepted for interface compatibility with ModelTrainer but unused.

  • unit_test_single_sample – Accepted for interface compatibility with ModelTrainer but unused.

  • log_frequency – Accepted for interface compatibility with ModelTrainer but unused.

  • gigabyte_vram – Accepted for interface compatibility with ModelTrainer but unused.

  • val_dataloader_weights – Accepted for interface compatibility with ModelTrainer but unused.

Variables:

  • test_dataloader (train_dataloader, val_dataloaders,) – Loaders created by _load_PEAL_dataset or _load_replication_dataset.

  • objective (torch.nn.Module) – Per-sample (reduction="none") loss selected from the task’s criterions.

Raises:

  • NotImplementedError – If replication_dataset is not None, "celeba" or "waterbirds".

  • Exception – If no supported architecture is found, or the task defines several non-l2 criteria or an unsupported one.

Methods

run

Train for training.max_epochs epochs with the GroupDRO loss.

run_epoch

Run one training or validation pass over loader.
__init__(adaptor_config='<PEAL_BASE>/configs/adaptors/test_dro.yaml', model_path=None, model=None, datasource=None, optimizer=None, criterions=None, logger=None, only_last_layer=False, unit_test_train_loop=False, unit_test_single_sample=False, log_frequency=1000, gigabyte_vram=None, val_dataloader_weights=[1.0])[source]

Unpack the config, build the dataloaders, the model and the objective.

Resolves model_path, loads either the PEAL dataset from datasource or one of the replication datasets (celeba, waterbirds), instantiates the architecture when model is None (optionally resetting its weights), and picks the single per-sample objective from task.criterions. See the class docstring for the parameters and the exceptions raised.

Parameters:

adaptor_config (dict | str | Path | AdaptorConfig)

run(continue_training=False, is_initialized=False)[source]

Train for training.max_epochs epochs with the GroupDRO loss.

Side effects: an existing model_path is renamed with an _old_ timestamp suffix (unless is_initialized), <model_path>/config.yaml is written, checkpoints/<epoch>.cpl is saved after every epoch and checkpoints/final.cpl whenever the validation avg_acc of the LossComputer improves.

Parameters:

  • continue_training (bool, optional) – Not supported; raises NotImplementedError if True.

  • is_initialized (bool, optional) – Skip moving an existing model directory out of the way.

Raises:

  • ValueError – If generalization_adjustment is neither a scalar nor a list of length 1 or n_groups.

  • NotImplementedError – If the training optimizer is not "sgd".

Notes

A ReduceLROnPlateau scheduler is created when config.scheduler is set but never stepped.

run_epoch(epoch, model, optimizer, loader, loss_computer, is_training, pbar=None)[source]

Run one training or validation pass over loader.

Parameters:

  • epoch (int) – Current epoch (only for bookkeeping).

  • model (torch.nn.Module) – Model to evaluate/update; switched to train or eval mode.

  • optimizer (torch.optim.Optimizer) – Stepped after every batch when is_training.

  • loader (DataLoader) – Yields (x, (y, confounder)) for PEAL datasets or (x, y, group) for replication datasets.

  • loss_computer (LossComputer) – Accumulates per-group statistics; its avg_actual_loss and avg_acc are printed at the end and then reset.

  • is_training (bool) – Whether gradients are enabled and the optimizer steps.

  • pbar (tqdm) – Progress bar advanced by one per batch; must not be None.