class_docstr: Implementation of the loss described in 'Whitening for
  Self-Supervised Representation Learning' [0].
  
  - [0] W-MSE, 2021, https://arxiv.org/pdf/2007.06346.pdf
  
  Examples:
      >>> # initialize loss function
      >>> loss_fn = WMSELoss(num_samples=2)
      >>> transform_fn = WMSETransform(num_samples=2)
      >>>
      >>> # generate the transformed samples
      >>> samples = transform_fn(image)
      >>>
      >>> # feed through encoder head
      >>> h = torch.cat([model(s) for s in samples])
      >>>
      >>> # calculate loss
      >>> loss = loss_fn(samples, h)
class_docstr_tok: ["Implementation","of","the","loss","described","in","'Whitening","for","Self-Supervised","Representation","Learning","'","[","0","]",".","-","[","0","]","W-MSE",",","2021",",","https",":","\/\/arxiv.org\/pdf\/2007.06346.pdf","Examples",":",">",">",">","#","initialize","loss","function",">",">",">","loss_fn","=","WMSELoss","(","num_samples=2",")",">",">",">","transform_fn","=","WMSETransform","(","num_samples=2",")",">",">",">",">",">",">","#","generate","the","transformed","samples",">",">",">","samples","=","transform_fn","(","image",")",">",">",">",">",">",">","#","feed","through","encoder","head",">",">",">","h","=","torch.cat","(","[","model","(","s",")","for","s","in","samples","]",")",">",">",">",">",">",">","#","calculate","loss",">",">",">","loss","=","loss_fn","(","samples",",","h",")"]
docstr: Parameters as described in [0]
  
  Args:
      embedding_dim:
          Dimensionality of the embedding.
      momentum:
          Momentum for the running statistics.
      eps:
          Epsilon for the running statistics.
      track_running_stats:
          Whether to track running statistics.
      w_iter:
          Number of iterations for the whitening.
      w_size:
          Sub-batch size to use for whitening.
      loss_fn:
          Loss function to use for the whitening.
      num_samples:
          Number of samples generated by the transforms for each image.
docstr_tok: ["Parameters","as","described","in","[","0","]","Args",":","embedding_dim",":","Dimensionality","of","the","embedding",".","momentum",":","Momentum","for","the","running","statistics",".","eps",":","Epsilon","for","the","running","statistics",".","track_running_stats",":","Whether","to","track","running","statistics",".","w_iter",":","Number","of","iterations","for","the","whitening",".","w_size",":","Sub-batch","size","to","use","for","whitening",".","loss_fn",":","Loss","function","to","use","for","the","whitening",".","num_samples",":","Number","of","samples","generated","by","the","transforms","for","each","image","."]
code: def __init__(
      self,
      embedding_dim: int = 128,
      momentum: float = 0.01,
      eps: float = 0.0,
      track_running_stats: bool = True,
      w_iter: int = 1,
      w_size: int = 256,
      loss_fn: Callable[
          [torch.Tensor, torch.Tensor], torch.Tensor
      ] = norm_mse_loss,
      num_samples: int = 2,
  ):
      """Parameters as described in [0]
  
      Args:
          embedding_dim:
              Dimensionality of the embedding.
          momentum:
              Momentum for the running statistics.
          eps:
              Epsilon for the running statistics.
          track_running_stats:
              Whether to track running statistics.
          w_iter:
              Number of iterations for the whitening.
          w_size:
              Sub-batch size to use for whitening.
          loss_fn:
              Loss function to use for the whitening.
          num_samples:
              Number of samples generated by the transforms for each image.
  
  
      """
      super().__init__()
      self.whitening = Whitening2d(
          num_features=embedding_dim,
          momentum=momentum,
          eps=eps,
          track_running_stats=track_running_stats,
      )
      if embedding_dim * 2 > w_size:
          raise ValueError(
              "w_size should be at least twice the size of embedding_dim to avoid instabiliy"
          )
      self.w_iter = w_iter
      self.w_size = w_size
      self.loss_f = loss_fn
      self.num_samples = num_samples
      self.num_pairs = num_samples * (num_samples - 1) // 2
  
file_path: lightly/lightly/loss/wmse_loss.py
imports: from typing import Callable
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new WMSELoss object from inherited base classes: torch.nn with example: obj = WMSELoss(embedding_dim, momentum, eps, track_running_stats, w_iter, w_size, loss_fn, num_samples)