class_docstr: Implementation of the loss function from MSN [0].
  
  Code inspired by [1].
  
  - [0]: Masked Siamese Networks, 2022, https://arxiv.org/abs/2204.07141
  - [1]: https://github.com/facebookresearch/msn
  
  Attributes:
      temperature:
          Similarities between anchors and targets are scaled by the inverse of
          the temperature. Must be in (0, inf).
      sinkhorn_iterations:
          Number of sinkhorn normalization iterations on the targets.
      regularization_weight:
          Weight factor lambda by which the regularization loss is scaled. Set to 0
          to disable regularization.
      me_max_weight:
          Deprecated, use `regularization_weight` instead. Takes precendence over
          `regularization_weight` if not None. Weight factor lambda by which the mean
          entropy maximization regularization loss is scaled. Set to 0 to disable
          mean entropy maximization reguliarization.
      gather_distributed:
          If True, then target probabilities are gathered from all GPUs.
  
   Examples:
  
      >>> # initialize loss function
      >>> loss_fn = MSNLoss()
      >>>
      >>> # generate anchors and targets of images
      >>> anchors = transforms(images)
      >>> targets = transforms(images)
      >>>
      >>> # feed through MSN model
      >>> anchors_out = model(anchors)
      >>> targets_out = model.target(targets)
      >>>
      >>> # calculate loss
      >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)
docstr: Computes the MSN loss for a set of anchors, targets and prototypes.
  
  Args:
      anchors:
          Tensor with shape (batch_size * anchor_views, dim).
      targets:
          Tensor with shape (batch_size, dim).
      prototypes:
          Tensor with shape (num_prototypes, dim).
      target_sharpen_temperature:
          Temperature used to sharpen the target probabilities.
  
  Returns:
      Mean loss over all anchors.
docstr_tok: ["Computes","the","MSN","loss","for","a","set","of","anchors",",","targets","and","prototypes",".","Args",":","anchors",":","Tensor","with","shape","(","batch_size","*","anchor_views",",","dim",")",".","targets",":","Tensor","with","shape","(","batch_size",",","dim",")",".","prototypes",":","Tensor","with","shape","(","num_prototypes",",","dim",")",".","target_sharpen_temperature",":","Temperature","used","to","sharpen","the","target","probabilities",".","Returns",":","Mean","loss","over","all","anchors","."]
code: def forward(
      self,
      anchors: Tensor,
      targets: Tensor,
      prototypes: Tensor,
      target_sharpen_temperature: float = 0.25,
  ) -> Tensor:
      """Computes the MSN loss for a set of anchors, targets and prototypes.
  
      Args:
          anchors:
              Tensor with shape (batch_size * anchor_views, dim).
          targets:
              Tensor with shape (batch_size, dim).
          prototypes:
              Tensor with shape (num_prototypes, dim).
          target_sharpen_temperature:
              Temperature used to sharpen the target probabilities.
  
      Returns:
          Mean loss over all anchors.
  
      """
      num_views = anchors.shape[0] // targets.shape[0]
      anchors = F.normalize(anchors, dim=1)
      targets = F.normalize(targets, dim=1)
      prototypes = F.normalize(prototypes, dim=1)
  
      # anchor predictions
      anchor_probs = prototype_probabilities(
          anchors, prototypes, temperature=self.temperature
      )
  
      # target predictions
      with torch.no_grad():
          target_probs = prototype_probabilities(
              targets, prototypes, temperature=self.temperature
          )
          target_probs = sharpen(
              target_probs, temperature=target_sharpen_temperature
          )
          if self.sinkhorn_iterations > 0:
              target_probs = sinkhorn(
                  probabilities=target_probs,
                  iterations=self.sinkhorn_iterations,
                  gather_distributed=self.gather_distributed,
              )
          target_probs = target_probs.repeat((num_views, 1))
  
      # cross entropy loss
      loss = torch.mean(
          torch.sum(torch.log(anchor_probs ** (-target_probs)), dim=1)
      )
  
      # regularization loss
      if self.regularization_weight > 0:
          mean_anchor_probs = torch.mean(anchor_probs, dim=0)
          reg_loss = self.regularization_loss(
              mean_anchor_probs=mean_anchor_probs
          )
          loss += self.regularization_weight * reg_loss
  
      return loss
  
file_path: lightly/lightly/loss/msn_loss.py
imports: import math
  import warnings
  from typing import Union
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
label_desc: Use image node_id 2 for calling the MSNLoss obj's underlying member method code with example usage: obj.forward(anchors, targets, prototypes, target_sharpen_temperature) and returns: loss