NEDSPretrain

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

Bases: core.trainer.BaseTrainer

Multimodal masked modelling for NEDS, over spikes and the behaviors alike.

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

log_training_results()

Log the training results, at the end of training.