(prototype) Move the NT-Xent equation into lightly.functional

lightly/functional/__init__.py

@@ -0,0 +1,8 @@
+"""Pure functions over tensors: one equation each, no parameters, no state.
+
+A block that owns parameters, buffers or optimiser state does not belong here.
+"""
+
+from lightly.functional.ntxent import ntxent
+
+__all__ = ["ntxent"]

lightly/functional/ntxent.py

@@ -0,0 +1,95 @@
+"""The NT-Xent equation, with no state and no distributed communication."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+from torch.nn.functional import cross_entropy
+
+__all__ = ["ntxent"]
+
+
+def ntxent(
+    z0: Tensor,
+    z1: Tensor,
+    *,
+    temperature: float,
+    negatives: Optional[Tensor] = None,
+    z0_all: Optional[Tensor] = None,
+    z1_all: Optional[Tensor] = None,
+    positives_at: int = 0,
+) -> Tensor:
+    """Normalized temperature-scaled cross entropy, as in SimCLR.
+
+    Every argument is a tensor or a number, so this function communicates with
+    no other rank and holds nothing between calls. Gathering the negatives is
+    :class:`~lightly.loss.ntx_ent_loss.NTXentLoss`'s job, and it passes the
+    result in through ``z0_all`` and ``z1_all``.
+
+    Reference:
+        SimCLR, 2020, https://arxiv.org/abs/2002.05709
+
+    Args:
+        z0:
+            Projections of the first view, shape ``(B, D)``, L2-normalized.
+        z1:
+            Projections of the second view, shape ``(B, D)``, L2-normalized.
+        temperature:
+            Scales the logits by its inverse.
+        negatives:
+            Negatives from a memory bank, shape ``(D, K)``. When given, the other
+            samples in the batch are not used as negatives.
+        z0_all:
+            The first view's projections across all ranks, shape ``(B * W, D)``.
+            Defaults to ``z0``, which is the single-rank case.
+        z1_all:
+            The second view's projections across all ranks. Defaults to ``z1``.
+        positives_at:
+            The column at which this rank's rows start inside ``z0_all``, so
+            ``rank * B``. Zero on one rank.
+
+    Returns:
+        The loss, as a scalar.
+    """
+    device = z0.device
+    batch_size = z0.size(0)
+
+    if negatives is not None:
+        # sim_pos[i] is the similarity of sample i to its positive pair, and
+        # sim_neg[i, j] its similarity to the j-th vector in the bank.
+        sim_pos = torch.einsum("nc,nc->n", z0, z1).unsqueeze(-1)
+        sim_neg = torch.einsum("nc,ck->nk", z0, negatives.to(device))
+        logits = torch.cat([sim_pos, sim_neg], dim=1) / temperature
+        labels = torch.zeros(logits.size(0), device=device, dtype=torch.long)
+        return cross_entropy(logits, labels)
+
+    z0_all = z0 if z0_all is None else z0_all
+    z1_all = z1 if z1_all is None else z1_all
+
+    # The similarities of a view with itself, which carry no signal and are
+    # dropped below. On one rank this is torch.eye.
+    rows = torch.arange(batch_size, device=device, dtype=torch.long)
+    diagonal = torch.zeros(batch_size, z0_all.size(0), device=device, dtype=torch.bool)
+    diagonal[rows, rows + positives_at] = True
+
+    # n is the local batch size and m the batch across ranks, so every block is
+    # (n, m).
+    logits_00 = torch.einsum("nc,mc->nm", z0, z0_all) / temperature
+    logits_01 = torch.einsum("nc,mc->nm", z0, z1_all) / temperature
+    logits_10 = torch.einsum("nc,mc->nm", z1, z0_all) / temperature
+    logits_11 = torch.einsum("nc,mc->nm", z1, z1_all) / temperature
+
+    logits_00 = logits_00[~diagonal].view(batch_size, -1)
+    logits_11 = logits_11[~diagonal].view(batch_size, -1)
+
+    logits = torch.cat(
+        [
+            torch.cat([logits_01, logits_00], dim=1),
+            torch.cat([logits_10, logits_11], dim=1),
+        ],
+        dim=0,
+    )
+    labels = (rows + positives_at).repeat(2)
+    return cross_entropy(logits, labels)

lightly/loss/__init__.py

@@ -9,6 +9,7 @@
 from lightly.loss.detcon_loss import DetConBLoss, DetConSLoss
 from lightly.loss.dino_loss import DINOLoss
 from lightly.loss.directclr_loss import DirectCLRLoss
+from lightly.loss.distributed import DistributedKind
 from lightly.loss.emp_ssl_loss import EMPSSLLoss
 from lightly.loss.frossl_loss import FroSSLLoss
 from lightly.loss.ibot_loss import IBOTPatchLoss, IBOTPlusPlusPatchLoss

lightly/loss/distributed.py

@@ -0,0 +1,35 @@
+"""How a loss behaves when the batch is split across ranks."""
+
+from enum import Enum, auto
+
+__all__ = ["DistributedKind"]
+
+
+class DistributedKind(Enum):
+    """What a loss needs from the ranks around it.
+
+    DDP averages the per-rank loss and the per-rank gradient. Whether that is
+    correct depends on the loss, and the difference is not visible in a forward
+    pass: issue #1920 shipped a correct SIGReg forward with a gradient off by
+    exactly ``1 / world_size``, and came back as #1977 against
+    ``BarlowTwinsLoss``. Declaring the kind is what selects the test.
+    """
+
+    RANK_LOCAL = auto()
+    """The per-rank value is a valid estimate, so DDP averaging is correct.
+
+    An MSE reconstruction loss.
+    """
+
+    GATHER_FOR_NEGATIVES = auto()
+    """Correct only if the features are gathered with gradient first.
+
+    NT-Xent and CLIP: every sample on every rank is a negative for every other.
+    """
+
+    GLOBAL_STATISTIC = auto()
+    """A function of the global batch, so averaging per rank is wrong.
+
+    SIGReg and the distribution-matching regularisers. The gap between the two
+    is the variance across ranks, which closes only once the ranks agree.
+    """

lightly/loss/ntx_ent_loss.py

@@ -9,6 +9,8 @@
 from torch import Tensor, nn
 from torch import distributed as torch_dist
 
+from lightly.functional.ntxent import ntxent
+from lightly.loss.distributed import DistributedKind
 from lightly.models.modules.memory_bank import MemoryBankModule
 from lightly.utils import dist
 
@@ -59,6 +61,8 @@ class NTXentLoss(nn.Module):
 
     """
 
+    distributed_kind = DistributedKind.GATHER_FOR_NEGATIVES
+
     def __init__(
         self,
         temperature: float = 0.5,
@@ -85,7 +89,6 @@ def __init__(
         )
         self.temperature = temperature
         self.gather_distributed = gather_distributed
-        self.cross_entropy = nn.CrossEntropyLoss(reduction="mean")
         self.eps = 1e-8
 
         if abs(self.temperature) < self.eps:
@@ -117,10 +120,8 @@ def forward(self, out0: Tensor, out1: Tensor) -> Tensor:
         Returns:
             Contrastive Cross Entropy Loss value.
         """
-        device = out0.device
-        batch_size, _ = out0.shape
-
-        # Normalize the output to length 1
+        # Normalize the output to length 1. The memory bank stores unit vectors,
+        # and ntxent takes them, so this happens here and only here.
         out0 = nn.functional.normalize(out0, dim=1)
         out1 = nn.functional.normalize(out1, dim=1)
 
@@ -132,65 +133,19 @@ def forward(self, out0: Tensor, out1: Tensor) -> Tensor:
         # negatives: shape: (embedding_size, memory_bank_size)
         out1, negatives = self.memory_bank.forward(out1, update=out0.requires_grad)
 
-        # Use cosine similarity (dot product) as all vectors are normalized to unit length
-        # Notation in einsum: n = batch_size, c = embedding_size and k = memory_bank_size.
-
         if negatives is not None:
-            # Use negatives from memory bank
-            negatives = negatives.to(device)
-
-            # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity
-            # of the i-th sample in the batch to its positive pair
-            sim_pos = torch.einsum("nc,nc->n", out0, out1).unsqueeze(-1)
-
-            # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity
-            # of the i-th sample to the j-th negative sample
-            sim_neg = torch.einsum("nc,ck->nk", out0, negatives)
-
-            # Set the labels to maximize sim_pos in relation to sim_neg
-            logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature
-            labels = torch.zeros(logits.shape[0], device=device, dtype=torch.long)
-
-        else:
-            # Use other samples from batch as negatives
-            # and create diagonal mask that only selects similarities between
-            # views of the same image
-            if self.gather_distributed and dist.world_size() > 1:
-                # Gather hidden representations from other processes
-                out0_large = torch.cat(dist.gather(out0), 0)
-                out1_large = torch.cat(dist.gather(out1), 0)
-                diag_mask = dist.eye_rank(batch_size, device=out0.device)
-            else:
-                # Single process
-                out0_large = out0
-                out1_large = out1
-                diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool)
-
-            # Calculate similiarities
-            # Here n = batch_size and m = batch_size * world_size
-            # The resulting vectors have shape (n, m)
-            logits_00 = torch.einsum("nc,mc->nm", out0, out0_large) / self.temperature
-            logits_01 = torch.einsum("nc,mc->nm", out0, out1_large) / self.temperature
-            logits_10 = torch.einsum("nc,mc->nm", out1, out0_large) / self.temperature
-            logits_11 = torch.einsum("nc,mc->nm", out1, out1_large) / self.temperature
-
-            # Remove simliarities between same views of the same image
-            logits_00 = logits_00[~diag_mask].view(batch_size, -1)
-            logits_11 = logits_11[~diag_mask].view(batch_size, -1)
-
-            # Concatenate logits
-            # The logits tensor in the end has shape (2*n, 2*m-1)
-            logits_0100 = torch.cat([logits_01, logits_00], dim=1)
-            logits_1011 = torch.cat([logits_10, logits_11], dim=1)
-            logits = torch.cat([logits_0100, logits_1011], dim=0)
-
-            # Create labels
-            labels = torch.arange(batch_size, device=device, dtype=torch.long)
-            if self.gather_distributed:
-                labels = labels + dist.rank() * batch_size
-            labels = labels.repeat(2)
-
-        # Calculate the cross-entropy loss
-        loss: Tensor = self.cross_entropy(logits, labels)
-
-        return loss
+            return ntxent(out0, out1, temperature=self.temperature, negatives=negatives)
+
+        if self.gather_distributed and dist.world_size() > 1:
+            # GATHER_FOR_NEGATIVES: every sample on every rank is a negative for
+            # every other, so the gather happens with gradient before the loss.
+            return ntxent(
+                out0,
+                out1,
+                temperature=self.temperature,
+                z0_all=torch.cat(dist.gather(out0), 0),
+                z1_all=torch.cat(dist.gather(out1), 0),
+                positives_at=dist.rank() * out0.size(0),
+            )
+
+        return ntxent(out0, out1, temperature=self.temperature)

tests/functional/test_ntxent.py

@@ -0,0 +1,69 @@
+"""Golden values for the NT-Xent equation.
+
+The numbers below are the second statement of the equation, as data rather than
+as code. Editing ``lightly/functional/ntxent.py`` moves them, and someone then
+has to approve the new number or restore the old one. A pinned loss sees changes
+above roughly 1e-5.
+"""
+
+from typing import Tuple
+
+import pytest
+import torch
+from torch import Tensor
+from torch.nn.functional import normalize
+
+from lightly.functional import ntxent
+from lightly.loss import NTXentLoss
+
+GOLDEN_LOSS = 5.350811958313  # seed 11, batch 8, dim 16, temperature 0.1
+GOLDEN_LOSS_MEMORY_BANK = 9.919633865356  # the same batch, 8 negatives from a bank
+TOLERANCE = 1e-5
+
+
+def batch(seed: int = 11) -> Tuple[Tensor, Tensor]:
+    torch.manual_seed(seed)
+    return (
+        normalize(torch.randn(8, 16), dim=1),
+        normalize(torch.randn(8, 16), dim=1),
+    )
+
+
+def test_golden_loss() -> None:
+    z0, z1 = batch()
+    assert float(ntxent(z0, z1, temperature=0.1)) == pytest.approx(
+        GOLDEN_LOSS, abs=TOLERANCE
+    )
+
+
+def test_golden_loss_with_a_memory_bank() -> None:
+    z0, z1 = batch()
+    torch.manual_seed(11)
+    negatives = normalize(torch.randn(8, 16), dim=1).T
+    loss = ntxent(z0, z1, temperature=0.1, negatives=negatives)
+    assert float(loss) == pytest.approx(GOLDEN_LOSS_MEMORY_BANK, abs=TOLERANCE)
+
+
+def test_the_module_computes_the_same_equation() -> None:
+    z0, z1 = batch()
+    module = NTXentLoss(temperature=0.1)
+    assert float(module(z0, z1)) == pytest.approx(GOLDEN_LOSS, abs=TOLERANCE)
+
+
+def test_temperature_scales_the_logits() -> None:
+    z0, z1 = batch()
+    assert float(ntxent(z0, z1, temperature=0.1)) > float(
+        ntxent(z0, z1, temperature=0.5)
+    )
+
+
+def test_positives_at_zero_is_the_single_rank_case() -> None:
+    z0, z1 = batch()
+    with_all = ntxent(z0, z1, temperature=0.1, z0_all=z0, z1_all=z1, positives_at=0)
+    assert torch.equal(with_all, ntxent(z0, z1, temperature=0.1))
+
+
+def test_the_loss_declares_how_it_behaves_across_ranks() -> None:
+    from lightly.loss import DistributedKind
+
+    assert NTXentLoss.distributed_kind is DistributedKind.GATHER_FOR_NEGATIVES