POYOPlusMultitaskPretrain

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

Bases: core.trainer.BaseTrainer

Supervised pretraining of POYO+ on every behavior at once, via its readout heads.

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.

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.

log_training_results()

Log the training results, at the end of training.