POSSMMultitaskPretrain

class pretrain.models.POSSMMultitaskPretrain(cfg, rank, world_size)

Bases: core.trainer.BaseTrainer

Multi-task pretrainer for POSSM.

Trains a single POSSM model on multiple decoding tasks simultaneously, using its MultitaskReadout. Per-task targets are gathered by the model’s input_fn (no per-sample dataset target extraction). Supports per-task eval metrics, configurable best-model selection (loss or metric), held-out test, and full ckpt resume.

setup(ckpt)

Set up the trainer.

This method is called after setting up distributed, loading the checkpoint, and setting up logging.

Parameters:

ckpt (dict | None) – The checkpoint dictionary, if any.

compute_multitask_loss(pred_dict, batch)

Aggregate per-task losses into a single training scalar.

Returns the unweighted task_losses for logging / metrics, plus a weighted sum sum_t (task_weights[t] * task_losses[t]) as the backward signal. Weighted-sum (not weighted-mean) is intentional: scaling a task’s weight up directly increases its gradient contribution, matching how POSSM’s recommended task_weights are usually written down.

train_epoch()

Perform a training epoch.

val_epoch()

Perform a validation epoch.

Called every val.every_n_epochs epochs, and at the end of training.

Returns:

bool, True if early stopping should be triggered, False (or falsey) otherwise.

Return type:

early_stop

test()

Perform a test epoch.

Called at the end of training if test.enable is True.

get_best_pbar_metrics()

Return the best validation metrics formatted for the epoch progress bar.

Returns:

Best validation metrics with the best/val/ prefix stripped,

or an empty dict if no validation has been run yet.

Return type:

dict

log_training_results()

Log the training results, at the end of training.