BYOL

class selfeeg.ssl.contrastive.BYOL(encoder: Module, projection_head: list[int] | Module, predictor: list[int] | Module = None, m: float = 0.999)[source]

Implementation of the BYOL SSL method.

To check how BYOL works, read the following paper [BYOL1] .

Parameters:

  • encoder (nn.Module) – The encoder part of the module. It is the one you wish to pretrain and transfer to the new model.

  • projection_head (Union[list[int], nn.Module]) –

    The projection head to use. It can be:

    1. an nn.Module

    2. a list of ints.

    In case a list is given, a nn.Sequential module with Dense, BatchNorm and Relu will be automtically created. The list will be used to set input and output dimension of each Dense Layer. For instance, if [64, 128, 64] is given, two hidden layers will be created. The first with input 64 and output 128, the second with input 128 and output 64.

  • predictor (Union[list[int], nn.Module]) – The predictor to put after the projection head. Accepted arguments are the same as for the projection_head.

  • m (float, optional) –

    The value of the momentum coefficient. Suggested values are in the range [0.9960, 0.9999].

    Default = 0.999

Warning

This class will not check the compatibility of the encoder’s output and the projection head’s input (as well as between the projection head and the predictor). Make sure that they have the same size.

References

[BYOL1]

J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al., “Bootstrap your own latent- a new approach to self-supervised learning,” Advances in neural information processing systems, vol. 33, pp. 21271– 21284, 2020.

Example

>>> import pickle, torch, selfeeg
>>> import selfeeg.dataloading as dl
>>> utils.create_dataset()
>>> def loadEEG(path, return_label=False):
...     with open(path, 'rb') as handle:
...         EEG = pickle.load(handle)
...     x , y= EEG['data'], EEG['label']
...     return (x, y) if return_label else x
>>> def loss_fineTuning(yhat, ytrue):
...     return F.binary_cross_entropy_with_logits(torch.squeeze(yhat), ytrue + 0.)
>>> torch.manual_seed(1234)
>>> # usual pipeline to construct the dataloader
>>> EEGlen = dl.get_eeg_partition_number('Simulated_EEG',freq=128, window=1,
...                              overlap=0.3, load_function=loadEEG)
>>> EEGsplit = dl.get_eeg_split_table (EEGlen, seed=1234)
>>> TrainSet = dl.EEGDataset(EEGlen,EEGsplit,[128,1,0.3],'train',False,loadEEG)
>>> TrainLoader = torch.utils.data.DataLoader(TrainSet, batch_size=32)
>>> enc = selfeeg.models.ShallowNetEncoder(8)
>>> byol = selfeeg.ssl.BYOL(enc, [16,32,32], [32,32])
>>> print( byol(torch.randn(32,8,128)).shape) # should return [32,32])
>>> loss_train = byol.fit(TrainLoader, 1, return_loss_info=True)
>>> print(loss_train[0][0]) # will return 2.16688
>>> loss_test = byol.test(TrainLoader)
>>> print(loss_test) # will return 1.2185
_update_momentum_encoder()[source]
fit(train_dataloader, epochs=1, optimizer=None, augmenter=None, loss_func: Callable = None, loss_args: list = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, verbose: bool = True, device: str = None, return_loss_info: bool = False)[source]

fit is a custom fit function designed to perform pretraining on a given model with the given dataloader.

Parameters:

  • train_dataloader (Dataloader) – The pytorch Dataloader used to get the training batches. It must return a batch as a single tensor X, thus without label tensor Y.

  • epochs (int, optional) –

    The number of training epochs. Must be an integer bigger than 0.

    Default = 1

  • optimizer (torch Optimizer, optional) –

    The optimizer used for weight’s update. It can be any optimizer provided in the torch.optim module. If not given Adam with default parameters will be instantiated.

    Default = torch.optim.Adam

  • augmenter (function, optional) –

    Any function (or callable object) used to perform data augmentation on the batch. It is highly suggested to resort to the augmentation module, which implements different data augmentation functions and classes to combine them. If none is given, a default augmentation with random vertical flip + random noise is applied. Note that in this case data augmentation is also performed on the validation set, since it is part of the SSL algorithm.

    Default = None

  • loss_func (Callable, optional) –

    The custom loss function. It can be any loss function which accepts as input only the model’s predictions (4 torch Tensor) as required arguments and loss_args as optional arguments. Check the input arguments of byol_loss to check how to design custom loss functions to give to this method.

    Default = None

  • loss_args (Union[list, dict], optional) –

    The optional arguments to pass to the function. It can be a list or a dict.

    Default = None

  • lr_scheduler (torch Scheduler) –

    A pytorch learning rate scheduler used to update the learning rate during the fine-tuning.

    Default = None

  • EarlyStopper (EarlyStopping, optional) –

    An instance of the provided EarlyStopping class.

    Default = None

  • validation_dataloader (Dataloader, optional) –

    The pytorch Dataloader used to get the validation batches. It must return a batch as a single tensor X, thus without label tensor Y. If not given, no validation loss will be calculated

    Default = None

  • verbose (bool, optional) –

    Whether to print a progression bar or not.

    Default = None

  • device (torch.device or str, optional) –

    The device to use for fine-tuning. If given as a string it will be converted in a torch.device instance. If not given, ‘cpu’ device will be used.

    Default = None

  • return_loss_info (bool, optional) –

    Whether to return the calculated training validation losses at each epoch.

    Default = False

Returns:

loss_info (dict, optional) – A dictionary with keys being the epoch number (as integer) and values a two element list with the average epoch’s training and validation loss.

Note

If an EarlyStopping instance is given with monitoring loss set to validation loss, but no validation dataloader is given, monitoring loss will be automatically set to training loss.

forward(x)[source]
test(test_dataloader, augmenter=None, loss_func: Callable = None, loss_args: list = [], verbose: bool = True, device: str = None)[source]

A method to evaluate the loss on a test dataloader. Parameters are the same as in the fit method, apart for the ones specific for the training which are removed.

It is rare to evaluate the pretraing loss function on a test set. Nevertheless this function provides a way to do that. An example of usage could be to assess the quality of the learned features on the fine-tuning dataset.