NuCLRPretrain#

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

Bases: core.trainer.BaseTrainer

Contrastive pretraining of the NuCLR unit encoder, against NuCLRLoss.

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]#

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