NDT2Pretrain#

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

Masked spike modelling for NDT2, which masks context tokens and accumulates gradients.

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.

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

get_best_pbar_metrics()[source]#

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()[source]#: Log the training results, at the end of training.