Repository: lightly-ai/lightly

Issue:
Add VICReg loss
It adds a VICReg loss that is described in this [paper](https://arxiv.org/abs/2105.04906) also as proposed in #779 
The implementation follows the official code, [here](https://github.com/facebookresearch/vicreg). In term of architecture, it's similar to BarlowTwins, so there is no need for a new projection head.

I would appreciate any feedback on this.


Discussion:

Resolving patch:
diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py
--- a/lightly/loss/__init__.py
+++ b/lightly/loss/__init__.py
@@ -11,3 +11,4 @@
 from lightly.loss.ntx_ent_loss import NTXentLoss
 from lightly.loss.swav_loss import SwaVLoss
 from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss
+from lightly.loss.vicreg_loss import VICRegLoss
diff --git a/lightly/loss/vicreg_loss.py b/lightly/loss/vicreg_loss.py
new file mode 100644
--- /dev/null
+++ b/lightly/loss/vicreg_loss.py
@@ -0,0 +1,102 @@
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from lightly.utils.dist import gather
+
+class VICRegLoss(torch.nn.Module):
+    """Implementation of the VICReg Loss from VICReg[0] paper.
+    This implementation follows the code published by the authors. [1]
+
+    [0] Bardes, A. et. al, 2022, VICReg... https://arxiv.org/abs/2105.04906
+    [1] https://github.com/facebookresearch/vicreg/
+        Examples:
+        >>> # initialize loss function
+        >>> loss_fn = VICRegLoss()
+        >>>
+        >>> # generate two random transforms of images
+        >>> t0 = transforms(images)
+        >>> t1 = transforms(images)
+        >>>
+        >>> # feed through model
+        >>> out0, out1 = model(t0, t1)
+        >>>
+        >>> # calculate loss
+        >>> loss = loss_fn(out0, out1)
+    """
+
+    def __init__(
+        self,
+        lambda_param: float = 25.0,
+        mu_param: float = 25.0,
+        nu_param: float = 1.0,
+        gather_distributed : bool = False,
+        eps = 0.0001
+    ):
+        """Lambda, mu and nu params configuration with default value like in [0]
+        Args:
+            lambda_param:
+                Coefficient for the invariance term of the loss
+                Defaults to 25.0 [0].
+            mu_param:
+                Coefficient for the variance term of the loss
+                Defaults to 25.0 [0].
+            nu_param:
+                Coefficient for the covariance term of the loss
+                Defaults to 1.0 [0].
+            gather_distributed:
+                If True then the cross-correlation matrices from all gpus are
+                gathered and summed before the loss calculation.
+            eps:
+                Numerical epsilon
+                Defaults to 0.0001 [1].
+        """
+        super(VICRegLoss, self).__init__()
+        self.lambda_param = lambda_param
+        self.mu_param = mu_param
+        self.nu_param = nu_param
+        self.gather_distributed = gather_distributed
+
+        self.eps = eps
+
+    def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
+        assert z_a.shape[0] > 1 and z_b.shape[0] > 1, f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and  {z_b.shape[0]}"
+        assert z_a.shape == z_b.shape, f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}."
+
+        # invariance term of the loss
+        repr_loss = F.mse_loss(z_a, z_b)
+
+        # gather all batches
+        if self.gather_distributed and dist.is_initialized():
+            world_size = dist.get_world_size()
+            if world_size > 1:
+                z_a = torch.cat(gather(z_a), dim=0)
+                z_b = torch.cat(gather(z_b), dim=0)
+
+        # normalize repr. along the batch dimension
+        z_a = z_a - z_a.mean(0) # NxD
+        z_b = z_b - z_b.mean(0) # NxD
+
+        N = z_a.size(0)
+        D = z_a.size(1)
+
+        # variance term of the loss
+        std_x = torch.sqrt(z_a.var(dim=0) + self.eps)
+        std_y = torch.sqrt(z_b.var(dim=0) + self.eps)
+        std_loss = 0.5 * (torch.mean(F.relu(1 - std_x)) + torch.mean(F.relu(1 - std_y)))
+
+        # covariance term of the loss
+        cov_x = (z_a.T @ z_a) / (N - 1)
+        cov_y = (z_b.T @ z_b) / (N - 1)
+
+        # compute off-diagonal elements
+        n, _ = cov_x.shape
+        off_diag_cov_x = cov_x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
+        off_diag_cov_y = cov_y.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
+
+        cov_loss = off_diag_cov_x.pow_(2).sum().div(D) + off_diag_cov_y.pow_(2).sum().div(D)
+
+        # loss
+        loss = self.lambda_param * repr_loss + self.mu_param * std_loss + self.nu_param * cov_loss
+
+        return loss


Test patch:
diff --git a/tests/loss/test_VICRegLoss.py b/tests/loss/test_VICRegLoss.py
new file mode 100644
--- /dev/null
+++ b/tests/loss/test_VICRegLoss.py
@@ -0,0 +1,42 @@
+import unittest
+import torch
+
+from lightly.loss import VICRegLoss
+
+class TestVICRegLoss(unittest.TestCase):
+    def test_forward_pass(self):
+        loss = VICRegLoss()
+        for bsz in range(2, 4):
+            x0 = torch.randn((bsz, 32))
+            x1 = torch.randn((bsz, 32))
+
+            # symmetry
+            l1 = loss(x0, x1)
+            l2 = loss(x1, x0)
+            self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available")
+    def test_forward_pass_cuda(self):
+        loss = VICRegLoss()
+        for bsz in range(2, 4):
+            x0 = torch.randn((bsz, 32)).cuda()
+            x1 = torch.randn((bsz, 32)).cuda()
+
+            # symmetry
+            l1 = loss(x0, x1)
+            l2 = loss(x1, x0)
+            self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0)
+
+    def test_forward_pass__error_batch_size_1(self):
+        loss = VICRegLoss()
+        x0 = torch.randn((1, 32))
+        x1 = torch.randn((1, 32))
+        with self.assertRaises(AssertionError):
+            loss(x0, x1)
+
+    def test_forward_pass__error_different_shapes(self):
+        loss = VICRegLoss()
+        x0 = torch.randn((2, 32))
+        x1 = torch.randn((2, 16))
+        with self.assertRaises(AssertionError):
+            loss(x0, x1)