From 7ff2c058955e9794a9260496601c7579d1960d8c Mon Sep 17 00:00:00 2001 From: KylevdLangemheen Date: Thu, 7 Aug 2025 11:44:35 +0200 Subject: [PATCH 1/3] Add first setup supcon loss --- lightly/loss/supcon_loss.py | 271 ++++++++++++++++++++++++++++++++++++ 1 file changed, 271 insertions(+) create mode 100644 lightly/loss/supcon_loss.py diff --git a/lightly/loss/supcon_loss.py b/lightly/loss/supcon_loss.py new file mode 100644 index 000000000..b9474ced2 --- /dev/null +++ b/lightly/loss/supcon_loss.py @@ -0,0 +1,271 @@ +""" Contrastive Loss Functions """ + +# Copyright (c) 2020. Lightly AG and its affiliates. +# All Rights Reserved + +from enum import Enum +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import Tensor +from torch import distributed as torch_dist +from torch import nn + +from lightly.utils import dist + + +def divide_no_nan(numerator: Tensor, denominator: Tensor) -> Tensor: + """Performs tensor division, setting result to zero where denominator is zero. + + Args: + numerator: + Numerator tensor. + denominator: + Denominator tensor with possible zeroes. + + Returns: + Result with zeros where denominator is zero. + """ + result = torch.zeros_like(numerator) + nonzero_mask = denominator != 0 + result[nonzero_mask] = numerator[nonzero_mask] / denominator[nonzero_mask] + return result + + +class ContrastMode(Enum): + """Contrast Mode Enum for SupCon Loss. + + Offers the three contrast modes as enum for the SupCon loss. The three modes are: + + - ContrastMode.ALL: Uses all positives and negatives. + - ContrastMode.ONE_POSITIVE: Uses only one positive, and all negatives. + - ContrastMode.ONLY_NEGATIVES: Uses no positives, only negatives. + """ + + ALL = 1 + ONE_POSITIVE = 2 + ONLY_NEGATIVES = 3 + + +class SupConLoss(nn.Module): + """Implementation of the Supervised Contrastive Loss. + + This implementation follows the SupCon[0] paper. + + - [0] SupCon, 2020, https://arxiv.org/abs/2004.11362 + + Attributes: + temperature: + Scale logits by the inverse of the temperature. + contrast_mode: + Whether to use all positives, one positive, or none. All negatives are + used in all cases. + gather_distributed: + If True then negatives from all GPUs are gathered before the + loss calculation. + + Raises: + ValueError: If abs(temperature) < 1e-8 to prevent divide by zero. + ValueError: If gather_distributed is True but torch.distributed is not available. + NotImplementedError: If contrast_mode is outside the accepted ContrastMode values. + + Examples: + >>> # initialize loss function without memory bank + >>> loss_fn = NTXentLoss(memory_bank_size=0) + >>> + >>> # generate two random transforms of images + >>> t0 = transforms(images) + >>> t1 = transforms(images) + >>> + >>> # feed through SimCLR or MoCo model + >>> out0, out1 = model(t0), model(t1) + >>> + >>> # calculate loss + >>> loss = loss_fn(out0, out1) + + """ + + def __init__( + self, + temperature: float = 0.5, + contrast_mode: ContrastMode = ContrastMode.ALL, + gather_distributed: bool = False, + ): + """Initializes the NTXentLoss module with the specified parameters. + + Args: + temperature: + Scale logits by the inverse of the temperature. + gather_distributed: + If True, negatives from all GPUs are gathered before the loss calculation. + + Raises: + ValueError: If temperature is less than 1e-8 to prevent divide by zero. + ValueError: If gather_distributed is True but torch.distributed is not available. + NotImplementedError: If contrast_mode is outside the accepted ContrastMode values. + """ + super().__init__() + self.temperature = temperature + self.contrast_mode = contrast_mode + self.positives_cap = -1 # Unused at the moment + self.gather_distributed = gather_distributed + self.cross_entropy = nn.CrossEntropyLoss(reduction="mean") + self.eps = 1e-8 + + if abs(self.temperature) < self.eps: + raise ValueError( + "Illegal temperature: abs({}) < 1e-8".format(self.temperature) + ) + if gather_distributed and not torch_dist.is_available(): + raise ValueError( + "gather_distributed is True but torch.distributed is not available. " + "Please set gather_distributed=False or install a torch version with " + "distributed support." + ) + + def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: + """Forward pass through Supervised Contrastive Loss. + + Computes the loss based on contrast_mode setting. + + Args: + features: + Tensor of at least 3 dimensions, corresponding to + (batch_size, num_views, ...) + labels: + Onehot labels for each sample. Must match shape + (batch_size, num_classes) + + Raises: + ValueError: If features does not have at least 3 dimensions. + ValueError: If number of labels does not match batch_size. + + Returns: + Supervised Contrastive Loss value. + """ + + device = features.device + batch_size, num_views = features.shape[:2] + + # Normalize the features to length 1 + features = F.normalize(features, dim=2) + + # Memory bank could be used here but labelled samples are not yet supported. + + # Use cosine similarity (dot product) as all vectors are normalized to unit length + + # Use other samples from different classes in batch as negatives + # and create diagonal mask that only selects similarities between + # views of the same image / same class + if self.gather_distributed and dist.world_size() > 1: + # Gather hidden representations and optional labels from other processes + global_features = torch.cat(dist.gather(features), 0) + diag_mask = dist.eye_rank(batch_size, device=device) + if labels is not None: + global_labels = torch.cat(dist.gather(labels), 0) + else: + # Single process + global_features = features + diag_mask = torch.eye(batch_size, device=device, dtype=torch.bool) + if labels is not None: + global_labels = labels + + # Use the diagonal mask if labels is none, else compute the mask based on labels + if labels is None: + # No labels, typical semi-supervised contrastive learning like SimCLR + mask = diag_mask + else: + mask = (labels @ global_labels.T).to(device) + + # Get features in shape [num_views * n, c] + all_global_features = global_features.permute(1, 0, 2).reshape( + -1, global_features.size(-1) + ) + + if self.contrast_mode == ContrastMode.ONE_POSITIVE: + anchor_features = features[:, 0] + num_anchor_views = 1 + else: + anchor_features = features.permute(1, 0, 2).reshape(-1, features.size(-1)) + num_anchor_views = num_views + + # Obtain the logits between anchor features and features across all processes + # Logits will be shaped [local_batch_size * num_anchor_views, global_batch_size * num_views] + # We then temperature scale it and subtract the max to improve numerical stability + logits = torch.einsum("nc,mc->nm", anchor_features, all_global_features) + logits /= self.temperature + logits -= logits.max(dim=1, keepdim=True)[0].detach() + exp_logits = torch.exp(logits) + + positives_mask, negatives_mask = self._create_tiled_masks( + mask, diag_mask, num_views, num_anchor_views, self.positives_cap + ) + num_positives_per_row = positives_mask.sum(dim=1) + + if self.contrast_mode == ContrastMode.ONE_POSITIVE: + denominator = exp_logits + (exp_logits * negatives_mask).sum( + dim=1, keepdim=True + ) + elif self.contrast_mode == ContrastMode.ALL: + denominator = (exp_logits * negatives_mask).sum(dim=1, keepdim=True) + denominator += (exp_logits * positives_mask).sum(dim=1, keepdim=True) + else: # ContrastMode.ONLY_NEGATIVES + denominator = (exp_logits * negatives_mask).sum(dim=1, keepdim=True) + + # num_positives_per_row can be zero iff 1 view is used. Here we use a safe + # dividing method seting those values to zero to prevent division by zero errors. + + # Only implements SupCon_{out} + log_probs = (logits - torch.log(denominator)) * positives_mask + log_probs = log_probs.sum(dim=1) + log_probs = divide_no_nan(log_probs, num_positives_per_row) + + loss = -log_probs + + if num_views != 1: + loss = loss.mean(dim=0) + else: + num_valid_views_per_sample = num_positives_per_row.unsqueeze(0) + loss = divide_no_nan(loss, num_valid_views_per_sample).squeeze() + + return loss + + def _create_tiled_masks( + self, untiled_mask, diagonal_mask, num_views, num_anchor_views, positives_cap + ) -> Tuple[Tensor, Tensor]: + # Get total batch size across all processes + print(untiled_mask.shape) + global_batch_size = untiled_mask.size(1) + + # Find index of the anchor for each sample + labels = torch.argmax(diagonal_mask.long(), dim=1) + + # Generate tiled labels across views + tiled_labels = [] + for i in range(num_anchor_views): + tiled_labels.append(labels + global_batch_size * i) + tiled_labels = torch.cat(tiled_labels, 0) + tiled_diagonal_mask = F.one_hot(tiled_labels, global_batch_size * num_views) + + # Mask to zero the diagonal at the end + all_but_diagonal_mask = 1 - tiled_diagonal_mask + + # All tiled positives + uncapped_positives_mask = torch.tile( + untiled_mask, [num_anchor_views, num_views] + ) + + # The negatives is simply the bitflipped positives + negatives_mask = 1 - uncapped_positives_mask + + # For when positives_cap is implemented + if positives_cap > -1: + raise NotImplementedError("Capping positives is not yet implemented.") + else: + positives_mask = uncapped_positives_mask + + # Zero out the self-contrast + positives_mask *= all_but_diagonal_mask + + return positives_mask, negatives_mask From b5cf5a0cedb6f6639f9557d677d8ffe818ec5f6b Mon Sep 17 00:00:00 2001 From: KylevdLangemheen Date: Thu, 7 Aug 2025 16:26:50 +0200 Subject: [PATCH 2/3] Start on supcon loss and tests (#1554) --- lightly/loss/__init__.py | 1 + lightly/loss/supcon_loss.py | 101 +++++++++++++++++++++++++++------ tests/loss/test_supcon_loss.py | 23 ++++++++ 3 files changed, 108 insertions(+), 17 deletions(-) create mode 100644 tests/loss/test_supcon_loss.py diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py index e0e2291cf..973459a3f 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -16,6 +16,7 @@ from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity from lightly.loss.ntx_ent_loss import NTXentLoss from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss +from lightly.loss.supcon_loss import SupConLoss from lightly.loss.swav_loss import SwaVLoss from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss from lightly.loss.tico_loss import TiCoLoss diff --git a/lightly/loss/supcon_loss.py b/lightly/loss/supcon_loss.py index b9474ced2..6dad9be83 100644 --- a/lightly/loss/supcon_loss.py +++ b/lightly/loss/supcon_loss.py @@ -48,6 +48,9 @@ class ContrastMode(Enum): ONLY_NEGATIVES = 3 +VALID_CONTRAST_MODES = set(item.name for item in ContrastMode) + + class SupConLoss(nn.Module): """Implementation of the Supervised Contrastive Loss. @@ -68,21 +71,24 @@ class SupConLoss(nn.Module): Raises: ValueError: If abs(temperature) < 1e-8 to prevent divide by zero. ValueError: If gather_distributed is True but torch.distributed is not available. - NotImplementedError: If contrast_mode is outside the accepted ContrastMode values. + ValueError: If contrast_mode is outside the accepted ContrastMode values. Examples: - >>> # initialize loss function without memory bank - >>> loss_fn = NTXentLoss(memory_bank_size=0) + >>> # initialize loss function + >>> loss_fn = SupConLoss() >>> - >>> # generate two random transforms of images + >>> # generate two or more views of images >>> t0 = transforms(images) >>> t1 = transforms(images) >>> - >>> # feed through SimCLR or MoCo model + >>> # feed through SimCLR model >>> out0, out1 = model(t0), model(t1) >>> + >>> # Stack views along 2nd dimensions + >>> features = torch.stack([out0, out1], dim=1) + >>> >>> # calculate loss - >>> loss = loss_fn(out0, out1) + >>> loss = loss_fn(features, labels) """ @@ -92,18 +98,21 @@ def __init__( contrast_mode: ContrastMode = ContrastMode.ALL, gather_distributed: bool = False, ): - """Initializes the NTXentLoss module with the specified parameters. + """Initializes the SupConLoss module with the specified parameters. Args: temperature: Scale logits by the inverse of the temperature. + contrast_mode: + Whether to use all positives, one positive, or none. All negatives are + used in all cases. gather_distributed: If True, negatives from all GPUs are gathered before the loss calculation. Raises: ValueError: If temperature is less than 1e-8 to prevent divide by zero. ValueError: If gather_distributed is True but torch.distributed is not available. - NotImplementedError: If contrast_mode is outside the accepted ContrastMode values. + ValueError: If contrast_mode is outside the accepted ContrastMode values. """ super().__init__() self.temperature = temperature @@ -124,6 +133,11 @@ def __init__( "distributed support." ) + if contrast_mode.name not in VALID_CONTRAST_MODES: + raise ValueError( + f"contrast_mode is {contrast_mode} but must be one of ContrastMode.{VALID_CONTRAST_MODES}" + ) + def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: """Forward pass through Supervised Contrastive Loss. @@ -140,14 +154,34 @@ def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: Raises: ValueError: If features does not have at least 3 dimensions. ValueError: If number of labels does not match batch_size. + ValueError: If labels is not one-hot encoded. Returns: Supervised Contrastive Loss value. """ + if len(features.shape) < 3: + raise ValueError( + f"Features must have at least 3 dimensions, got {len(features.shape)}." + ) + device = features.device batch_size, num_views = features.shape[:2] + if labels is not None and labels.size(0) != batch_size: + raise ValueError( + f"When setting labels, labels must match batch_size {batch_size}, got {labels.size(0)}." + ) + + if labels is not None: + if not self._is_one_hot(labels): + raise ValueError( + "labels must be a 2D matrix representing the one-hot encoded classes." + ) + + # Flatten the features in case they are still images or other + features = features.flatten(2) + # Normalize the features to length 1 features = F.normalize(features, dim=2) @@ -178,31 +212,43 @@ def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: else: mask = (labels @ global_labels.T).to(device) - # Get features in shape [num_views * n, c] + # Get features in shape [num_views * batch_size, c] all_global_features = global_features.permute(1, 0, 2).reshape( -1, global_features.size(-1) ) if self.contrast_mode == ContrastMode.ONE_POSITIVE: + # We take only the first view as anchor anchor_features = features[:, 0] num_anchor_views = 1 else: + # We take all views as anchors in the same shape as the global features anchor_features = features.permute(1, 0, 2).reshape(-1, features.size(-1)) num_anchor_views = num_views # Obtain the logits between anchor features and features across all processes # Logits will be shaped [local_batch_size * num_anchor_views, global_batch_size * num_views] # We then temperature scale it and subtract the max to improve numerical stability + # In the einsum, n is local_batch_size * num_anchor_views, m is global_batch_size * num_views, + # and c is the flattened feature length + # Note: features are ordered by view first, i.e. first all samples of view 0, then all samples + # of view 1, and so on. logits = torch.einsum("nc,mc->nm", anchor_features, all_global_features) logits /= self.temperature logits -= logits.max(dim=1, keepdim=True)[0].detach() exp_logits = torch.exp(logits) + # Get the positive and negative masks for numerator & denominator positives_mask, negatives_mask = self._create_tiled_masks( - mask, diag_mask, num_views, num_anchor_views, self.positives_cap + mask.long(), + diag_mask.long(), + num_views, + num_anchor_views, + self.positives_cap, ) num_positives_per_row = positives_mask.sum(dim=1) + # Calculate denominator based on contrast_mode if self.contrast_mode == ContrastMode.ONE_POSITIVE: denominator = exp_logits + (exp_logits * negatives_mask).sum( dim=1, keepdim=True @@ -216,13 +262,14 @@ def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: # num_positives_per_row can be zero iff 1 view is used. Here we use a safe # dividing method seting those values to zero to prevent division by zero errors. - # Only implements SupCon_{out} + # Only implements SupCon_{out}. log_probs = (logits - torch.log(denominator)) * positives_mask log_probs = log_probs.sum(dim=1) log_probs = divide_no_nan(log_probs, num_positives_per_row) loss = -log_probs + # Adjust for num_positives_per_row being zero when using exactly 1 view if num_views != 1: loss = loss.mean(dim=0) else: @@ -232,21 +279,27 @@ def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: return loss def _create_tiled_masks( - self, untiled_mask, diagonal_mask, num_views, num_anchor_views, positives_cap + self, + untiled_mask: Tensor, + diagonal_mask: Tensor, + num_views: int, + num_anchor_views: int, + positives_cap: int, ) -> Tuple[Tensor, Tensor]: # Get total batch size across all processes - print(untiled_mask.shape) global_batch_size = untiled_mask.size(1) # Find index of the anchor for each sample - labels = torch.argmax(diagonal_mask.long(), dim=1) + labels = torch.argmax(diagonal_mask, dim=1) # Generate tiled labels across views tiled_labels = [] for i in range(num_anchor_views): tiled_labels.append(labels + global_batch_size * i) - tiled_labels = torch.cat(tiled_labels, 0) - tiled_diagonal_mask = F.one_hot(tiled_labels, global_batch_size * num_views) + tiled_labels_tensor = torch.cat(tiled_labels, 0) + tiled_diagonal_mask = F.one_hot( + tiled_labels_tensor, global_batch_size * num_views + ) # Mask to zero the diagonal at the end all_but_diagonal_mask = 1 - tiled_diagonal_mask @@ -257,7 +310,7 @@ def _create_tiled_masks( ) # The negatives is simply the bitflipped positives - negatives_mask = 1 - uncapped_positives_mask + negatives_mask = 1.0 - uncapped_positives_mask # For when positives_cap is implemented if positives_cap > -1: @@ -269,3 +322,17 @@ def _create_tiled_masks( positives_mask *= all_but_diagonal_mask return positives_mask, negatives_mask + + def _is_one_hot(self, tensor: Tensor) -> bool: + # Tensor is not a 2D matrix + if tensor.ndim != 2: + return False + + # Check values are only 0 or 1 + is_binary = ((tensor == 0) | (tensor == 1)).all() + + # Check each row sums to 1 + row_sums = tensor.sum(dim=1) + has_single_one = (row_sums == 1).all() + + return bool(is_binary.item() and has_single_one.item()) diff --git a/tests/loss/test_supcon_loss.py b/tests/loss/test_supcon_loss.py new file mode 100644 index 000000000..e8859ddde --- /dev/null +++ b/tests/loss/test_supcon_loss.py @@ -0,0 +1,23 @@ +import pytest +import torch +from torch import Tensor + +from lightly.loss import NTXentLoss, SupConLoss + + +class TestSupConLoss: + def test_simple_input(self) -> None: + my_input = torch.rand([3, 2, 4]) + my_label = Tensor([[1, 0], [0, 1], [0, 1]]) + my_loss = SupConLoss() + my_loss(my_input, my_label) + + def test_unsup_equal_to_simclr(self) -> None: + supcon = SupConLoss(temperature=0.5, gather_distributed=False) + ntxent = NTXentLoss( + temperature=0.5, memory_bank_size=0, gather_distributed=False + ) + features = torch.rand((8, 2, 10)) + supcon_loss = supcon(features) + ntxent_loss = ntxent(features[:, 0, :], features[:, 1, :]) + assert (supcon_loss - ntxent_loss).pow(2).item() == pytest.approx(0.0) From f7d94fb06a64e1d87ce9d1dbdeb1eabfc3203eb6 Mon Sep 17 00:00:00 2001 From: KylevdLangemheen Date: Mon, 25 Aug 2025 19:29:27 +0200 Subject: [PATCH 3/3] Simplify SupConLoss (#1877) --- lightly/loss/supcon_loss.py | 312 ++++++++------------------------- tests/loss/test_supcon_loss.py | 124 +++++++++++-- 2 files changed, 185 insertions(+), 251 deletions(-) diff --git a/lightly/loss/supcon_loss.py b/lightly/loss/supcon_loss.py index 6dad9be83..26d48d892 100644 --- a/lightly/loss/supcon_loss.py +++ b/lightly/loss/supcon_loss.py @@ -3,8 +3,7 @@ # Copyright (c) 2020. Lightly AG and its affiliates. # All Rights Reserved -from enum import Enum -from typing import Optional, Tuple +from typing import Optional import torch import torch.nn.functional as F @@ -15,42 +14,6 @@ from lightly.utils import dist -def divide_no_nan(numerator: Tensor, denominator: Tensor) -> Tensor: - """Performs tensor division, setting result to zero where denominator is zero. - - Args: - numerator: - Numerator tensor. - denominator: - Denominator tensor with possible zeroes. - - Returns: - Result with zeros where denominator is zero. - """ - result = torch.zeros_like(numerator) - nonzero_mask = denominator != 0 - result[nonzero_mask] = numerator[nonzero_mask] / denominator[nonzero_mask] - return result - - -class ContrastMode(Enum): - """Contrast Mode Enum for SupCon Loss. - - Offers the three contrast modes as enum for the SupCon loss. The three modes are: - - - ContrastMode.ALL: Uses all positives and negatives. - - ContrastMode.ONE_POSITIVE: Uses only one positive, and all negatives. - - ContrastMode.ONLY_NEGATIVES: Uses no positives, only negatives. - """ - - ALL = 1 - ONE_POSITIVE = 2 - ONLY_NEGATIVES = 3 - - -VALID_CONTRAST_MODES = set(item.name for item in ContrastMode) - - class SupConLoss(nn.Module): """Implementation of the Supervised Contrastive Loss. @@ -61,64 +24,55 @@ class SupConLoss(nn.Module): Attributes: temperature: Scale logits by the inverse of the temperature. - contrast_mode: - Whether to use all positives, one positive, or none. All negatives are - used in all cases. gather_distributed: If True then negatives from all GPUs are gathered before the - loss calculation. - + loss calculation. If a memory bank is used and gather_distributed is True, + then tensors from all gpus are gathered before the memory bank is updated. + rescale: + Optionally rescale final loss by the temperature for stability. Raises: ValueError: If abs(temperature) < 1e-8 to prevent divide by zero. - ValueError: If gather_distributed is True but torch.distributed is not available. - ValueError: If contrast_mode is outside the accepted ContrastMode values. Examples: - >>> # initialize loss function - >>> loss_fn = SupConLoss() + >>> # initialize loss function without memory bank + >>> loss_fn = NTXentLoss(memory_bank_size=0) >>> - >>> # generate two or more views of images + >>> # generate two random transforms of images >>> t0 = transforms(images) >>> t1 = transforms(images) >>> - >>> # feed through SimCLR model + >>> # feed through SimCLR or MoCo model >>> out0, out1 = model(t0), model(t1) >>> - >>> # Stack views along 2nd dimensions - >>> features = torch.stack([out0, out1], dim=1) - >>> >>> # calculate loss - >>> loss = loss_fn(features, labels) + >>> loss = loss_fn(out0, out1) """ def __init__( self, temperature: float = 0.5, - contrast_mode: ContrastMode = ContrastMode.ALL, gather_distributed: bool = False, + rescale: bool = True, ): """Initializes the SupConLoss module with the specified parameters. Args: temperature: Scale logits by the inverse of the temperature. - contrast_mode: - Whether to use all positives, one positive, or none. All negatives are - used in all cases. gather_distributed: If True, negatives from all GPUs are gathered before the loss calculation. + rescale: + Optionally rescale final loss by the temperature for stability. Raises: ValueError: If temperature is less than 1e-8 to prevent divide by zero. ValueError: If gather_distributed is True but torch.distributed is not available. - ValueError: If contrast_mode is outside the accepted ContrastMode values. """ super().__init__() self.temperature = temperature - self.contrast_mode = contrast_mode - self.positives_cap = -1 # Unused at the moment self.gather_distributed = gather_distributed + self.rescale = rescale self.cross_entropy = nn.CrossEntropyLoss(reduction="mean") self.eps = 1e-8 @@ -133,206 +87,82 @@ def __init__( "distributed support." ) - if contrast_mode.name not in VALID_CONTRAST_MODES: - raise ValueError( - f"contrast_mode is {contrast_mode} but must be one of ContrastMode.{VALID_CONTRAST_MODES}" - ) - - def forward(self, features: Tensor, labels: Optional[Tensor] = None) -> Tensor: + def forward( + self, out0: Tensor, out1: Tensor, labels: Optional[Tensor] = None + ) -> Tensor: """Forward pass through Supervised Contrastive Loss. Computes the loss based on contrast_mode setting. Args: - features: - Tensor of at least 3 dimensions, corresponding to - (batch_size, num_views, ...) + out0: + Output projections of the first set of transformed images. + Shape: (batch_size, embedding_size) + out1: + Output projections of the second set of transformed images. + Shape: (batch_size, embedding_size) labels: - Onehot labels for each sample. Must match shape - (batch_size, num_classes) - - Raises: - ValueError: If features does not have at least 3 dimensions. - ValueError: If number of labels does not match batch_size. - ValueError: If labels is not one-hot encoded. + Onehot labels for each sample. Must be a vector of length `batch_size`. Returns: Supervised Contrastive Loss value. """ + # Stack the views for efficient computation + # Allows for more views to be added easily + features = (out0, out1) + n_views = len(features) + out_small = torch.vstack(features) - if len(features.shape) < 3: - raise ValueError( - f"Features must have at least 3 dimensions, got {len(features.shape)}." - ) - - device = features.device - batch_size, num_views = features.shape[:2] - - if labels is not None and labels.size(0) != batch_size: - raise ValueError( - f"When setting labels, labels must match batch_size {batch_size}, got {labels.size(0)}." - ) - - if labels is not None: - if not self._is_one_hot(labels): - raise ValueError( - "labels must be a 2D matrix representing the one-hot encoded classes." - ) - - # Flatten the features in case they are still images or other - features = features.flatten(2) - - # Normalize the features to length 1 - features = F.normalize(features, dim=2) - - # Memory bank could be used here but labelled samples are not yet supported. + device = out_small.device + batch_size = out_small.shape[0] // n_views - # Use cosine similarity (dot product) as all vectors are normalized to unit length + # Normalize the output to length 1 + out_small = nn.functional.normalize(out_small, dim=1) - # Use other samples from different classes in batch as negatives - # and create diagonal mask that only selects similarities between - # views of the same image / same class + # Gather hidden representations from other processes if distributed + # and compute the diagonal self-contrast mask if self.gather_distributed and dist.world_size() > 1: - # Gather hidden representations and optional labels from other processes - global_features = torch.cat(dist.gather(features), 0) - diag_mask = dist.eye_rank(batch_size, device=device) - if labels is not None: - global_labels = torch.cat(dist.gather(labels), 0) + out_large = torch.cat(dist.gather(out_small), 0) + diag_mask = dist.eye_rank(n_views * batch_size, device=device) else: # Single process - global_features = features - diag_mask = torch.eye(batch_size, device=device, dtype=torch.bool) - if labels is not None: - global_labels = labels - - # Use the diagonal mask if labels is none, else compute the mask based on labels - if labels is None: - # No labels, typical semi-supervised contrastive learning like SimCLR - mask = diag_mask - else: - mask = (labels @ global_labels.T).to(device) + out_large = out_small + diag_mask = torch.eye(n_views * batch_size, device=device, dtype=torch.bool) - # Get features in shape [num_views * batch_size, c] - all_global_features = global_features.permute(1, 0, 2).reshape( - -1, global_features.size(-1) - ) - - if self.contrast_mode == ContrastMode.ONE_POSITIVE: - # We take only the first view as anchor - anchor_features = features[:, 0] - num_anchor_views = 1 - else: - # We take all views as anchors in the same shape as the global features - anchor_features = features.permute(1, 0, 2).reshape(-1, features.size(-1)) - num_anchor_views = num_views - - # Obtain the logits between anchor features and features across all processes - # Logits will be shaped [local_batch_size * num_anchor_views, global_batch_size * num_views] - # We then temperature scale it and subtract the max to improve numerical stability - # In the einsum, n is local_batch_size * num_anchor_views, m is global_batch_size * num_views, - # and c is the flattened feature length - # Note: features are ordered by view first, i.e. first all samples of view 0, then all samples - # of view 1, and so on. - logits = torch.einsum("nc,mc->nm", anchor_features, all_global_features) + # Use cosine similarity (dot product) as all vectors are normalized to unit length + # Calculate similiarities + logits = out_small @ out_large.T logits /= self.temperature - logits -= logits.max(dim=1, keepdim=True)[0].detach() - exp_logits = torch.exp(logits) - # Get the positive and negative masks for numerator & denominator - positives_mask, negatives_mask = self._create_tiled_masks( - mask.long(), - diag_mask.long(), - num_views, - num_anchor_views, - self.positives_cap, - ) - num_positives_per_row = positives_mask.sum(dim=1) + # Set self-similarities to infinitely small value + logits[diag_mask] = -1e9 - # Calculate denominator based on contrast_mode - if self.contrast_mode == ContrastMode.ONE_POSITIVE: - denominator = exp_logits + (exp_logits * negatives_mask).sum( - dim=1, keepdim=True - ) - elif self.contrast_mode == ContrastMode.ALL: - denominator = (exp_logits * negatives_mask).sum(dim=1, keepdim=True) - denominator += (exp_logits * positives_mask).sum(dim=1, keepdim=True) - else: # ContrastMode.ONLY_NEGATIVES - denominator = (exp_logits * negatives_mask).sum(dim=1, keepdim=True) - - # num_positives_per_row can be zero iff 1 view is used. Here we use a safe - # dividing method seting those values to zero to prevent division by zero errors. - - # Only implements SupCon_{out}. - log_probs = (logits - torch.log(denominator)) * positives_mask - log_probs = log_probs.sum(dim=1) - log_probs = divide_no_nan(log_probs, num_positives_per_row) - - loss = -log_probs - - # Adjust for num_positives_per_row being zero when using exactly 1 view - if num_views != 1: - loss = loss.mean(dim=0) - else: - num_valid_views_per_sample = num_positives_per_row.unsqueeze(0) - loss = divide_no_nan(loss, num_valid_views_per_sample).squeeze() + # Create labels if None + if labels is None: + labels = torch.arange(batch_size, device=device, dtype=torch.long) + if self.gather_distributed: + labels = labels + dist.rank() * batch_size + labels = labels.repeat(n_views) + + # Soft labels are 0 unless the logit represents a similarity + # between two of the same classes. We manually set self-similarity + # (same view of the same item) to 0. When not 0, the value is + # 1 / n, where n is the number of positive samples + # (different views of the same item, and all views of other items sharing + # classes with the item) + soft_labels = torch.eq(labels, labels.view(-1, 1)).float() + soft_labels.fill_diagonal_(0.0) + soft_labels /= soft_labels.sum(dim=1) + + # Compute log probabilities + log_proba = F.log_softmax(logits, dim=-1) + + # Compute soft cross-entropy loss + loss = (soft_labels * log_proba).sum(-1) + loss = -loss.mean() + + # Optional: rescale for stable training + if self.rescale: + loss *= self.temperature return loss - - def _create_tiled_masks( - self, - untiled_mask: Tensor, - diagonal_mask: Tensor, - num_views: int, - num_anchor_views: int, - positives_cap: int, - ) -> Tuple[Tensor, Tensor]: - # Get total batch size across all processes - global_batch_size = untiled_mask.size(1) - - # Find index of the anchor for each sample - labels = torch.argmax(diagonal_mask, dim=1) - - # Generate tiled labels across views - tiled_labels = [] - for i in range(num_anchor_views): - tiled_labels.append(labels + global_batch_size * i) - tiled_labels_tensor = torch.cat(tiled_labels, 0) - tiled_diagonal_mask = F.one_hot( - tiled_labels_tensor, global_batch_size * num_views - ) - - # Mask to zero the diagonal at the end - all_but_diagonal_mask = 1 - tiled_diagonal_mask - - # All tiled positives - uncapped_positives_mask = torch.tile( - untiled_mask, [num_anchor_views, num_views] - ) - - # The negatives is simply the bitflipped positives - negatives_mask = 1.0 - uncapped_positives_mask - - # For when positives_cap is implemented - if positives_cap > -1: - raise NotImplementedError("Capping positives is not yet implemented.") - else: - positives_mask = uncapped_positives_mask - - # Zero out the self-contrast - positives_mask *= all_but_diagonal_mask - - return positives_mask, negatives_mask - - def _is_one_hot(self, tensor: Tensor) -> bool: - # Tensor is not a 2D matrix - if tensor.ndim != 2: - return False - - # Check values are only 0 or 1 - is_binary = ((tensor == 0) | (tensor == 1)).all() - - # Check each row sums to 1 - row_sums = tensor.sum(dim=1) - has_single_one = (row_sums == 1).all() - - return bool(is_binary.item() and has_single_one.item()) diff --git a/tests/loss/test_supcon_loss.py b/tests/loss/test_supcon_loss.py index e8859ddde..2d1496fda 100644 --- a/tests/loss/test_supcon_loss.py +++ b/tests/loss/test_supcon_loss.py @@ -1,23 +1,127 @@ +from typing import List + import pytest import torch +from pytest_mock import MockerFixture from torch import Tensor +from torch import distributed as dist +from torch import nn from lightly.loss import NTXentLoss, SupConLoss class TestSupConLoss: + temperature = 0.5 + + def test__gather_distributed(self, mocker: MockerFixture) -> None: + mock_is_available = mocker.patch.object(dist, "is_available", return_value=True) + SupConLoss(gather_distributed=True) + mock_is_available.assert_called_once() + + def test__gather_distributed_dist_not_available( + self, mocker: MockerFixture + ) -> None: + mock_is_available = mocker.patch.object( + dist, "is_available", return_value=False + ) + with pytest.raises(ValueError): + SupConLoss(gather_distributed=True) + mock_is_available.assert_called_once() + def test_simple_input(self) -> None: - my_input = torch.rand([3, 2, 4]) - my_label = Tensor([[1, 0], [0, 1], [0, 1]]) + out1 = torch.rand((3, 10)) + out2 = torch.rand((3, 10)) + my_label = Tensor([0, 1, 1]) my_loss = SupConLoss() - my_loss(my_input, my_label) + my_loss(out1, out2, my_label) def test_unsup_equal_to_simclr(self) -> None: - supcon = SupConLoss(temperature=0.5, gather_distributed=False) - ntxent = NTXentLoss( - temperature=0.5, memory_bank_size=0, gather_distributed=False - ) - features = torch.rand((8, 2, 10)) - supcon_loss = supcon(features) - ntxent_loss = ntxent(features[:, 0, :], features[:, 1, :]) + supcon = SupConLoss(temperature=self.temperature, rescale=False) + ntxent = NTXentLoss(temperature=self.temperature) + out1 = torch.rand((8, 10)) + out2 = torch.rand((8, 10)) + supcon_loss = supcon(out1, out2) + ntxent_loss = ntxent(out1, out2) assert (supcon_loss - ntxent_loss).pow(2).item() == pytest.approx(0.0) + + @pytest.mark.parametrize("labels", [[0, 0, 0, 0], [0, 1, 1, 1], [0, 1, 2, 3]]) + def test_equivalence(self, labels: List[int]) -> None: + DistributedSupCon = SupConLoss(temperature=self.temperature) + NonDistributedSupCon = SupConLossNonDistributed(temperature=self.temperature) + out1 = nn.functional.normalize(torch.rand(4, 10), dim=-1) + out2 = nn.functional.normalize(torch.rand(4, 10), dim=-1) + test_labels = Tensor(labels) + + loss1 = DistributedSupCon(out1, out2, test_labels) + loss2 = NonDistributedSupCon( + torch.vstack((out1, out2)), test_labels.view(-1, 1) + ) + + assert (loss1 - loss2).pow(2).item() == pytest.approx(0.0) + + +class SupConLossNonDistributed(nn.Module): + def __init__( + self, + temperature: float = 0.1, + ): + """Contrastive Learning Loss Function: SupConLoss and InfoNCE Loss. Non-distributed version by Yutong. + + SupCon from Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. + InfoNCE (NT-Xent) from SimCLR: https://arxiv.org/pdf/2002.05709.pdf. + + Adapted from Yonglong Tian's work at https://github.com/HobbitLong/SupContrast/blob/master/losses.py and + https://github.com/google-research/syn-rep-learn/blob/main/StableRep/models/losses.py. + + The function first creates a contrastive mask of shape [batch_size * n_views, batch_size * n_views], where + mask_{i,j}=1 if sample j has the same class as sample i, except for the sample i itself. + + Next, it computes the logits from the features and then computes the soft cross-entropy loss. + + The loss is rescaled by the temperature parameter. + + For self-supervised learning, the labels should be the indices of the samples. In this case it is equivalent to InfoNCE loss. + + Attributes: + - temperature (float): A temperature parameter to control the similarity. Default is 0.1. + + Args: + - features (torch.Tensor): hidden vector of shape [batch_size * n_views, ...]. + - labels (torch.Tensor): ground truth of shape [batch_size, 1]. + """ + super().__init__() + + self.temperature = temperature + + def forward( + self, + features: torch.Tensor, + labels: torch.Tensor, + ) -> torch.Tensor: + # create n-viewed mask + labels_n_views = labels.contiguous().repeat( + features.shape[0] // labels.shape[0], 1 + ) # [batch_size * n_views, 1] + contrastive_mask_n_views = torch.eq( + labels_n_views, labels_n_views.T + ).float() # [batch_size * n_views, batch_size * n_views] + contrastive_mask_n_views.fill_diagonal_(0) # mask-out self-contrast cases + + # compute logits + logits = ( + torch.matmul(features, features.T) / self.temperature + ) # [batch_size * n_views, batch_size * n_views] + logits.fill_diagonal_(-1e9) # suppress logit for self-contrast cases + + # compute log probabilities and soft labels + soft_label = contrastive_mask_n_views / contrastive_mask_n_views.sum(dim=1) + log_proba = nn.functional.log_softmax(logits, dim=-1) + + # compute soft cross-entropy loss + loss_all = torch.sum(soft_label * log_proba, dim=-1) + loss = -loss_all.mean() + + # rescale for stable training + loss *= self.temperature + + return loss