class_docstr: Implementation of the CO2 regularizer [0] for self-supervised learning.
  
  [0] CO2, 2021, https://arxiv.org/abs/2010.02217
  
  Attributes:
      alpha:
          Weight of the regularization term.
      t_consistency:
          Temperature used during softmax calculations.
      memory_bank_size:
          Size of the memory bank as (num_features, dim) tuple. num_features is the
          number of negatives stored in the bank. If set to 0, the memory bank is
          disabled. Deprecated: If only a single integer is passed, it is interpreted
          as the number of features and the feature dimension is inferred from the
          first batch stored in the memory bank. Leaving out the feature dimension
          might lead to errors in distributed training.
  
  Examples:
      >>> # initialize loss function for MoCo
      >>> loss_fn = NTXentLoss(memory_bank_size=(4096, 128))
      >>>
      >>> # initialize CO2 regularizer
      >>> co2 = CO2Regularizer(alpha=1.0, memory_bank_size=(4096, 128))
      >>>
      >>> # generate two random trasnforms of images
      >>> t0 = transforms(images)
      >>> t1 = transforms(images)
      >>>
      >>> # feed through the MoCo model
      >>> out0, out1 = model(t0, t1)
      >>>
      >>> # calculate loss and apply regularizer
      >>> loss = loss_fn(out0, out1) + co2(out0, out1)
docstr: Computes the soft pseudo labels across negative samples.
  
  Args:
      out0:
          Output projections of the first set of transformed images (query).
          Shape: bsz x n_ftrs
      out1:
          Output projections of the second set of transformed images (positive sample).
          Shape: bsz x n_ftrs
      negatives:
          Negative samples to compare against. If this is None, the second
          batch of images will be used as negative samples.
          Shape: memory_bank_size x n_ftrs
  
  Returns:
      Log probability that a positive samples will classify each negative
      sample as the positive sample.
      Shape: bsz x (bsz - 1) or bsz x memory_bank_size
docstr_tok: ["Computes","the","soft","pseudo","labels","across","negative","samples",".","Args",":","out0",":","Output","projections","of","the","first","set","of","transformed","images","(","query",")",".","Shape",":","bsz","x","n_ftrs","out1",":","Output","projections","of","the","second","set","of","transformed","images","(","positive","sample",")",".","Shape",":","bsz","x","n_ftrs","negatives",":","Negative","samples","to","compare","against",".","If","this","is","None",",","the","second","batch","of","images","will","be","used","as","negative","samples",".","Shape",":","memory_bank_size","x","n_ftrs","Returns",":","Log","probability","that","a","positive","samples","will","classify","each","negative","sample","as","the","positive","sample",".","Shape",":","bsz","x","(","bsz","-","1",")","or","bsz","x","memory_bank_size"]
code: def _get_pseudo_labels(
      self,
      out0: torch.Tensor,
      out1: torch.Tensor,
      negatives: torch.Tensor = None,
  ):
      """Computes the soft pseudo labels across negative samples.
  
      Args:
          out0:
              Output projections of the first set of transformed images (query).
              Shape: bsz x n_ftrs
          out1:
              Output projections of the second set of transformed images (positive sample).
              Shape: bsz x n_ftrs
          negatives:
              Negative samples to compare against. If this is None, the second
              batch of images will be used as negative samples.
              Shape: memory_bank_size x n_ftrs
  
      Returns:
          Log probability that a positive samples will classify each negative
          sample as the positive sample.
          Shape: bsz x (bsz - 1) or bsz x memory_bank_size
  
      """
      batch_size, _ = out0.shape
      if negatives is None:
          # use second batch as negative samples
          # l_pos has shape bsz x 1 and l_neg has shape bsz x bsz
          l_pos = torch.einsum("nc,nc->n", [out0, out1]).unsqueeze(-1)
          l_neg = torch.einsum("nc,ck->nk", [out0, out1.t()])
          # remove elements on the diagonal
          # l_neg has shape bsz x (bsz - 1)
          l_neg = l_neg.masked_select(
              ~torch.eye(batch_size, dtype=bool, device=l_neg.device)
          ).view(batch_size, batch_size - 1)
      else:
          # use memory bank as negative samples
          # l_pos has shape bsz x 1 and l_neg has shape bsz x memory_bank_size
          negatives = negatives.to(out0.device)
          l_pos = torch.einsum("nc,nc->n", [out0, out1]).unsqueeze(-1)
          l_neg = torch.einsum(
              "nc,ck->nk", [out0, negatives.clone().detach()]
          )
  
      # concatenate such that positive samples are at index 0
      logits = torch.cat([l_pos, l_neg], dim=1)
      # divide by temperature
      logits = logits / self.t_consistency
  
      # the input to kl_div is expected to be log(p)
      return torch.nn.functional.log_softmax(logits, dim=-1)
  
file_path: lightly/lightly/loss/regularizer/co2.py
imports: from typing import Sequence, Union
  import torch
  from lightly.models.modules.memory_bank import MemoryBankModule
label_desc: Use image node_id 2 for calling the CO2Regularizer obj's underlying member method code with example usage: obj._get_pseudo_labels(out0, out1, negatives) and returns: torch