Repository: lightly-ai/lightly

Issue:
BarlowTwins



Discussion:

Resolving patch:
diff --git a/benchmarks/imagenet/resnet50/barlowtwins.py b/benchmarks/imagenet/resnet50/barlowtwins.py
new file mode 100644
--- /dev/null
+++ b/benchmarks/imagenet/resnet50/barlowtwins.py
@@ -0,0 +1,112 @@
+import copy
+from typing import List, Tuple
+
+import torch
+from pytorch_lightning import LightningModule
+from torch import Tensor
+from torch.nn import Identity
+from torchvision.models import resnet50
+
+from lightly.loss import BarlowTwinsLoss
+from lightly.models.modules import BarlowTwinsProjectionHead
+from lightly.models.utils import get_weight_decay_parameters
+from lightly.transforms import BYOLTransform
+from lightly.utils.benchmarking import OnlineLinearClassifier
+from lightly.utils.lars import LARS
+from lightly.utils.scheduler import CosineWarmupScheduler
+
+
+class BarlowTwins(LightningModule):
+    def __init__(self, batch_size_per_device: int, num_classes: int) -> None:
+        super().__init__()
+        self.save_hyperparameters()
+        self.batch_size_per_device = batch_size_per_device
+
+        resnet = resnet50()
+        resnet.fc = Identity()  # Ignore classification head
+        self.backbone = resnet
+        self.projection_head = BarlowTwinsProjectionHead()
+        self.criterion = BarlowTwinsLoss(lambda_param=5e-3, gather_distributed=True)
+
+        self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.backbone(x)
+
+    def training_step(
+        self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int
+    ) -> Tensor:
+        # Forward pass and loss calculation.
+        views, targets = batch[0], batch[1]
+        features = self.forward(torch.cat(views)).flatten(start_dim=1)
+        z = self.projection_head(features)
+        z0, z1 = z.chunk(len(views))
+        loss = self.criterion(z0, z1)
+
+        self.log(
+            "train_loss", loss, prog_bar=True, sync_dist=True, batch_size=len(targets)
+        )
+
+        # Online linear evaluation.
+        cls_loss, cls_log = self.online_classifier.training_step(
+            (features.detach(), targets.repeat(len(views))), batch_idx
+        )
+        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))
+        return loss + cls_loss
+
+    def validation_step(
+        self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int
+    ) -> Tensor:
+        images, targets = batch[0], batch[1]
+        features = self.forward(images).flatten(start_dim=1)
+        cls_loss, cls_log = self.online_classifier.validation_step(
+            (features.detach(), targets), batch_idx
+        )
+        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))
+        return cls_loss
+
+    def configure_optimizers(self):
+        lr_factor = self.batch_size_per_device * self.trainer.world_size / 256
+
+        # Don't use weight decay for batch norm, bias parameters, and classification
+        # head to improve performance.
+        params, params_no_weight_decay = get_weight_decay_parameters(
+            [self.backbone, self.projection_head]
+        )
+        optimizer = LARS(
+            [
+                {"name": "barlowtwins", "params": params},
+                {
+                    "name": "barlowtwins_no_weight_decay",
+                    "params": params_no_weight_decay,
+                    "weight_decay": 0.0,
+                    "lr": 0.0048 * lr_factor,
+                },
+                {
+                    "name": "online_classifier",
+                    "params": self.online_classifier.parameters(),
+                    "weight_decay": 0.0,
+                },
+            ],
+            lr=0.2 * lr_factor,
+            momentum=0.9,
+            weight_decay=1.5e-6,
+        )
+
+        scheduler = {
+            "scheduler": CosineWarmupScheduler(
+                optimizer=optimizer,
+                warmup_epochs=(
+                    self.trainer.estimated_stepping_batches
+                    / self.trainer.max_epochs
+                    * 10
+                ),
+                max_epochs=self.trainer.estimated_stepping_batches,
+            ),
+            "interval": "step",
+        }
+        return [optimizer], [scheduler]
+
+
+# BarlowTwins uses same transform as BYOL.
+transform = BYOLTransform()
diff --git a/benchmarks/imagenet/resnet50/main.py b/benchmarks/imagenet/resnet50/main.py
--- a/benchmarks/imagenet/resnet50/main.py
+++ b/benchmarks/imagenet/resnet50/main.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Sequence, Union
 
+import barlowtwins
 import byol
 import dcl
 import dclw
@@ -46,6 +47,10 @@
 parser.add_argument("--skip-finetune-eval", action="store_true")
 
 METHODS = {
+    "barlowtwins": {
+        "model": barlowtwins.BarlowTwins,
+        "transform": barlowtwins.transform,
+    },
     "byol": {"model": byol.BYOL, "transform": byol.transform},
     "dcl": {"model": dcl.DCL, "transform": dcl.transform},
     "dclw": {"model": dclw.DCLW, "transform": dclw.transform},
diff --git a/lightly/loss/barlow_twins_loss.py b/lightly/loss/barlow_twins_loss.py
--- a/lightly/loss/barlow_twins_loss.py
+++ b/lightly/loss/barlow_twins_loss.py
@@ -1,5 +1,8 @@
+from typing import Tuple
+
 import torch
 import torch.distributed as dist
+import torch.nn.functional as F
 
 
 class BarlowTwinsLoss(torch.nn.Module):
@@ -48,17 +51,14 @@ def __init__(self, lambda_param: float = 5e-3, gather_distributed: bool = False)
             )
 
     def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
-        device = z_a.device
-
         # normalize repr. along the batch dimension
-        z_a_norm = (z_a - z_a.mean(0)) / z_a.std(0)  # NxD
-        z_b_norm = (z_b - z_b.mean(0)) / z_b.std(0)  # NxD
+        z_a_norm, z_b_norm = _normalize(z_a, z_b)
 
         N = z_a.size(0)
-        D = z_a.size(1)
 
         # cross-correlation matrix
-        c = torch.mm(z_a_norm.T, z_b_norm) / N  # DxD
+        c = z_a_norm.T @ z_b_norm
+        c.div_(N)
 
         # sum cross-correlation matrix between multiple gpus
         if self.gather_distributed and dist.is_initialized():
@@ -67,10 +67,31 @@ def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
                 c = c / world_size
                 dist.all_reduce(c)
 
-        # loss
-        c_diff = (c - torch.eye(D, device=device)).pow(2)  # DxD
-        # multiply off-diagonal elems of c_diff by lambda
-        c_diff[~torch.eye(D, dtype=bool)] *= self.lambda_param
-        loss = c_diff.sum()
+        invariance_loss = torch.diagonal(c).add_(-1).pow_(2).sum()
+        redundancy_reduction_loss = _off_diagonal(c).pow_(2).sum()
+        loss = invariance_loss + self.lambda_param * redundancy_reduction_loss
 
         return loss
+
+
+def _normalize(
+    z_a: torch.Tensor, z_b: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Helper function to normalize tensors along the batch dimension."""
+    combined = torch.stack([z_a, z_b], dim=0)  # Shape: 2 x N x D
+    normalized = F.batch_norm(
+        combined.flatten(0, 1),
+        running_mean=None,
+        running_var=None,
+        weight=None,
+        bias=None,
+        training=True,
+    ).view_as(combined)
+    return normalized[0], normalized[1]
+
+
+def _off_diagonal(x):
+    # return a flattened view of the off-diagonal elements of a square matrix
+    n, m = x.shape
+    assert n == m
+    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()


Test patch:
diff --git a/tests/loss/test_barlow_twins_loss.py b/tests/loss/test_barlow_twins_loss.py
--- a/tests/loss/test_barlow_twins_loss.py
+++ b/tests/loss/test_barlow_twins_loss.py
@@ -1,10 +1,51 @@
 import pytest
+import torch
+import torch.nn as nn
 from pytest_mock import MockerFixture
 from torch import distributed as dist
 
 from lightly.loss.barlow_twins_loss import BarlowTwinsLoss
 
 
+class BarlowTwinsLossReference(torch.nn.Module):
+    def __init__(
+        self,
+        projector_dim: int = 8196,
+        lambda_param: float = 5e-3,
+        gather_distributed: bool = False,
+    ):
+        super(BarlowTwinsLossReference, self).__init__()
+        # normalization layer for the representations z1 and z2
+        self.bn = nn.BatchNorm1d(projector_dim, affine=False)
+        self.lambda_param = lambda_param
+        self.gather_distributed = gather_distributed
+
+    def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
+        # code from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
+
+        N = z_a.size(0)
+
+        # empirical cross-correlation matrix
+        c = self.bn(z_a).T @ self.bn(z_b)
+
+        # sum the cross-correlation matrix between all gpus
+        c.div_(N)
+        if self.gather_distributed:
+            torch.distributed.all_reduce(c)
+
+        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
+        off_diag = off_diagonal(c).pow_(2).sum()
+        loss = on_diag + self.lambda_param * off_diag
+        return loss
+
+
+def off_diagonal(x):
+    # return a flattened view of the off-diagonal elements of a square matrix
+    n, m = x.shape
+    assert n == m
+    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
+
+
 class TestBarlowTwinsLoss:
     def test__gather_distributed(self, mocker: MockerFixture) -> None:
         mock_is_available = mocker.patch.object(dist, "is_available", return_value=True)
@@ -20,3 +61,25 @@ def test__gather_distributed_dist_not_available(
         with pytest.raises(ValueError):
             BarlowTwinsLoss(gather_distributed=True)
         mock_is_available.assert_called_once()
+
+    def test__loss_matches_reference_loss(self) -> None:
+        batch_size = 32
+        projector_dim = 8196
+        lambda_param = 5e-3
+        gather_distributed = False
+        loss = BarlowTwinsLoss(
+            lambda_param=lambda_param, gather_distributed=gather_distributed
+        )
+        loss_ref = BarlowTwinsLossReference(
+            projector_dim=projector_dim,
+            lambda_param=lambda_param,
+            gather_distributed=gather_distributed,
+        )
+
+        z_a = torch.randn(batch_size, projector_dim)
+        z_b = torch.randn(batch_size, projector_dim)
+
+        loss_out = loss(z_a, z_b)
+        loss_ref_out = loss_ref(z_a, z_b)
+
+        assert torch.allclose(loss_out, loss_ref_out, rtol=1e-3, atol=1e-3)