peal.training.trainers.ModelTrainer

class peal.training.trainers.ModelTrainer[source]

Bases: object

Train a predictor from a PredictorConfig with early stopping.

On construction the model, dataloaders (via create_dataloaders_from_datasource), optimizer (sgd with a 0.95 exponential schedule, adam or adamw), criterions and Logger are built from the config. training.class_balanced wraps the train loader in a DataloaderMixer and splits every validation loader into one class-restricted copy per class. Training happens in fit().

Parameters:

  • config (str, dict or PredictorConfig) – Loaded with load_yaml_config; consumes model_path, only_last_layer, tracking_level and the training, task, data and architecture sections.

  • model_path (str, optional) – Output directory; overrides config.model_path.

  • model (torch.nn.Module, optional) – Pre-built model; otherwise get_predictor() is used.

  • datasource (optional) – Datasets or dataloaders handed to create_dataloaders_from_datasource.

  • optimizer (torch.optim.Optimizer, optional) – Overrides the optimizer built from training.optimizer.

  • criterions (dict, optional) – Overrides the criterions built from task.criterions.

  • logger (Logger, optional) – Overrides the default Logger.

  • only_last_layer (bool, optional) – Train only the final linear layer (model.fc or the last weight and bias). Defaults to config.only_last_layer.

  • unit_test_train_loop (bool, optional) – Stop every epoch after two batches.

  • unit_test_single_sample (bool, optional) – Replace every batch by the logger’s fixed test sample.

  • log_frequency (int, optional) – Stored on the instance; not used by the loop itself.

  • val_dataloader_weights (list of float, optional) – Weights for averaging validation accuracies across loaders.

Variables:

  • regularization_level (float) – Multiplier for the l1/l2/orthogonality criterions; raised by fit() when overfitting is detected.

  • attacker (PGD_L2) – Only present when training.adv_training is set.

Methods

fit

Run _fit() with every parameter the optimizer does not update frozen.

run_epoch

Run one pass over dataloader in train or validation mode.
__init__(config, model_path=None, model=None, datasource=None, optimizer=None, criterions=None, logger=None, only_last_layer=None, unit_test_train_loop=False, unit_test_single_sample=False, log_frequency=1000, val_dataloader_weights=[1.0])[source]

Build model, dataloaders, optimizer, criterions and logger.

See the class docstring for the parameters.

run_epoch(dataloader, mode='train', pbar=None, debug=False)[source]

Run one pass over dataloader in train or validation mode.

Per batch: optional diffusion augmentation, label smoothing, mixup and PGD adversarial attack (train mode only), forward pass (with latents when the lc criterion is configured), the weighted sum of task.criterions (regularisers scaled by regularization_level), loss.backward() and, in train mode, an optimizer step. Every step is reported to self.logger and the progress bar; the loop stops after training.steps_per_epoch batches when that is set.

Parameters:

  • dataloader (torch.utils.data.DataLoader or DataloaderMixer) – Batches of (X, y); when the loader has return_src the source distribution is tracked for the progress bar.

  • mode (str, optional) – "train" or "validation_<i>"; passed to the logger.

  • pbar (tqdm) – Progress bar with a stored_values dict (required).

  • debug (bool, optional) – Unused.

Returns:

(last_batch_loss, accuracy) where the accuracy comes from logger.log_epoch; the loss is 0.0 when no batch was yielded.

Return type:

tuple

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

Run _fit() with every parameter the optimizer does not update frozen.

only_last_layer hands the optimizer just the head, but the backbone kept requires_grad=True, so every step stored its activations and computed gradients that were thrown away: a DINOv3 ViT-L linear probe needed a full backward pass (and ran out of memory next to other jobs on 2026-09-28). The flags are restored afterwards because the same model object can later be fine-tuned in full (CFKD), where permanently frozen parameters would silently stop training.

Parameters:

  • continue_training (bool, optional) – Passed to _fit().

  • is_initialized (bool, optional) – Passed to _fit().