VICReg

class selfeeg.ssl.contrastive.VICReg(encoder: Module, projection_head: list[int] | Module)[source]

Implementation of the VICReg SSL method.

To check how VICReg works, read the following paper [VIC1] .

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.

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

[VIC1]

A. Bardes, J. Ponce, and Y. LeCun, “Vicreg: Variance- invariance-covariance regularization for self-supervised learning,” arXiv preprint arXiv:2105.04906, 2021.

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)
>>> vic = selfeeg.ssl.VICReg(enc, [16,32,32])
>>> print( vic(torch.randn(32,8,128)).shape) # should return [32,32])
>>> loss_train = vic.fit(TrainLoader, 2, return_loss_info=True)
>>> print(loss_train[0][0]) # will return 21.6086
>>> loss_test = vic.test(TrainLoader)
>>> print(loss_test) # will return 21.6785
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, cat_augmentations: bool = False, 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 as required arguments and loss_args as optional arguments. If not given VICReg loss will be automatically used. Check the input arguments of vicreg_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

  • cat_augmentations (bool, optional) –

    Whether to calculate the loss on the cat version of the two projection’s or not. It might affect some statistical layers.

    Default = False

  • 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.

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

Evaluate the loss on a test dataloader. Parameters are the same as described in the fit method, aside for those related to model 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.