RRRSingleTaskPretrain#

class pretrain.models.RRRSingleTaskPretrain(cfg, rank, world_size)[source]#

Bases: core.trainer.BaseTrainer

Supervised pretraining of the shared basis V on one behavior.

setup(ckpt)[source]#

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()[source]#

Perform a training epoch.

val_epoch()[source]#

Score pooled over sessions; best.pt is what the probe loads.

log_training_results()[source]#

Log the training results, at the end of training.