peal.adaptors.group_dro_v2.GroupDROv2

class peal.adaptors.group_dro_v2.GroupDROv2

Bases: Adaptor

Adaptor that retrains a classifier with the group-robust objective.

Parameters:

adaptor_config (GroupDROv2Config) – See GroupDROv2Config.

Variables:

• original_model (torch.nn.Module) – The loaded classifier, kept untouched.

• model (torch.nn.Module) – The deep copy that is actually retrained.

• test_group_counts (train_group_counts, val_group_counts,) – Samples per group, counted once by compute_group_sizes().

• test_data_unpoisoned (DataLoader or None) – Test loader built from config.unpoisoned_data.

Methods

run	Entry point of the adaptor interface; runs train().
run_epoch	Run one pass over loader, training or evaluating.
train	Retrain the model with the group-robust objective and evaluate it.

__init__(adaptor_config)

Prepare the run directory, load the model and count the group sizes.

Seeds torch, moves an existing base_dir aside with a timestamp suffix, writes config.yaml into the fresh one, loads the classifier, builds the train/validation dataloaders with groups enabled and dict-style batches, and optionally the unpoisoned test loader. Counting the group sizes walks each loader once.

Parameters:

adaptor_config (GroupDROv2Config)

run()

Entry point of the adaptor interface; runs train().

train()

Retrain the model with the group-robust objective and evaluate it.

Runs one validation (and, if configured, test) pass at epoch -1 as a baseline, then training.max_epochs train/validation rounds with SGD (momentum 0.9, config.weight_decay). The epoch’s validation score is the worst or the average group accuracy depending on training.early_stopping_goal; a new best writes model.cpl into base_dir. With automatic_adjustment the generalization adjustment is re-estimated each epoch as the train/validation gap scaled by sqrt(group_counts). Finally the best checkpoint is scored on the unpoisoned test split and its average and worst group accuracy are logged as test_final_*.

Side effects: TensorBoard logs under <base_dir>/logs, per-epoch checkpoints under <base_dir>/checkpoints when config.save_intermediate.

Raises:

NotImplementedError – If training.early_stopping_goal is neither "worst_group_accuracy" nor "average_group_accuracy".

run_epoch(epoch, model, optimizer, loader, loss_computer, log_writer=None, mode='train')

Run one pass over loader, training or evaluating.

Each batch is expected as a dict with x, y and has_confounder; the group index is 2 * y + has_confounder. Gradients are only enabled, and the optimizer only steps, in "train" mode. After the pass the loss computer’s statistics are written to TensorBoard under the <mode>_ prefix, and reset for the next training epoch (the evaluation computers are short-lived and discarded by the caller instead).

Parameters:

• epoch (int) – Step the statistics are logged at; -1 for the baseline pass.

• model (torch.nn.Module) – Model to train or evaluate.

• optimizer (torch.optim.Optimizer) – Only used in "train" mode.

• loader (torch.utils.data.DataLoader) – Loader yielding dict batches with group information.

• loss_computer (GroupDroLossComputer) – Computes the loss and accumulates the statistics.

• log_writer (torch.utils.tensorboard.SummaryWriter, optional) – When None, nothing is logged.

• mode (str, optional) – "train", "val" or "test".