simclr_loss

selfeeg.losses.losses.simclr_loss(projections: Tensor, projections_norm: bool = True, temperature: float = 0.15) → Tensor[source]

simclr_loss computes the normalized temperature-scaled cross entropy loss [NTXent] , which is used in many contrastive learning algorithm. It is basically a simple implementation of the InfoNCE_loss provided in the official simCLR repository [simgit] using only torch functions.

Parameters:

  • projections (torch.Tensor) – 2D Tensor where projections[0:N/2] are the projections of one batch augmented version and projections[N/2:] are the projections of the other batch augmented version

  • projections_norm (bool, optional) –

    Whether to normalize the projections or not.

    Default = True

  • temperature (float, optional) –

    Temperature coefficient of the NTX_ent loss (See references to check loss formula).

    Default = 0.15

Returns:

loss (torch.Tensor) – The calculated loss.

Note

Looking at some implementations (e.g. the one in lightlyAI), the returned loss seems to be double. However the function add_contrastive_loss in the original repo returns the same value as this implementation.

References

[NTXent]

Chen et al. A Simple Framework for Contrastive Learning of Visual Representations. (2020). https://doi.org/10.48550/arXiv.2002.05709

[simgit]

To check the original tensorflow implementation visit the following repository: https://github.com/google-research/simclr (look at the function add_contrastive_loss in objective.py)

Example

>>> import torch
>>> import selfeeg.losses
>>> torch.manual_seed(1234)
>>> projections = torch.randn(64, 32)
>>> loss = losses.simclr_loss(projections)
>>> print(loss) # will return 10.2866