class_docstr: The criterion corresponding to the Barlow Twins loss as defined in the paper
  https://arxiv.org/abs/2103.03230v1.
  
  Args:
      lambda_ (float):        weight on the off-diagonal terms. It controls the trade-off
                              between the importance given to the invariance term versus the
                              redundancy reduction term.
      scale_loss (float):     In order to match the code that was used to develop Barlow
                              Twins, we include an additional parameter, scale-loss, that
                              multiplies the loss by a constant factor. We are working on a
                              version that will not require this parameter.
      embedding_dim (int):    dimensionality of the representation
class_docstr_tok: ["The","criterion","corresponding","to","the","Barlow","Twins","loss","as","defined","in","the","paper","https",":","\/\/arxiv.org\/abs\/2103.03230v1",".","Args",":","lambda_","(","float",")",":","weight","on","the","off-diagonal","terms",".","It","controls","the","trade-off","between","the","importance","given","to","the","invariance","term","versus","the","redundancy","reduction","term",".","scale_loss","(","float",")",":","In","order","to","match","the","code","that","was","used","to","develop","Barlow","Twins",",","we","include","an","additional","parameter",",","scale-loss",",","that","multiplies","the","loss","by","a","constant","factor",".","We","are","working","on","a","version","that","will","not","require","this","parameter",".","embedding_dim","(","int",")",":","dimensionality","of","the","representation"]
docstr: Calculate the loss. Operates on embeddings tensor.
  
  Args:
      embedding (torch.Tensor):   NxEMBEDDING_DIM
                                  Must contain the concatenated embeddings
                                  of the two image copies:
                                  [emb_img1_0, emb_img2_0, ....., emb_img1_1, emb_img2_1,...]
docstr_tok: ["Calculate","the","loss",".","Operates","on","embeddings","tensor",".","Args",":","embedding","(","torch.Tensor",")",":","NxEMBEDDING_DIM","Must","contain","the","concatenated","embeddings","of","the","two","image","copies",":","[","emb_img1_0",",","emb_img2_0",",",".....",",","emb_img1_1",",","emb_img2_1",",","...","]"]
code: def forward(self, embedding: torch.Tensor) -> torch.tensor:
      """
      Calculate the loss. Operates on embeddings tensor.
  
      Args:
          embedding (torch.Tensor):   NxEMBEDDING_DIM
                                      Must contain the concatenated embeddings
                                      of the two image copies:
                                      [emb_img1_0, emb_img2_0, ....., emb_img1_1, emb_img2_1,...]
      """
      assert embedding.ndim == 2 and embedding.shape[1] == int(
          self.embedding_dim
      ), f"Incorrect embedding shape: {embedding.shape} but expected Nx{self.embedding_dim}"
  
      batch_size = embedding.shape[0]
      assert (
          batch_size % self.num_copies == 0
      ), f"Batch size {batch_size} should be divisible by num_copies ({self.num_copies})."
  
      # normalize embeddings along the batch dimension
      embedding_normed = SyncNormalizeFunction.apply(
          embedding, self.eps
      )
  
      # split embedding between copies
      embedding_normed_a, embedding_normed_b = torch.split(
          embedding_normed,
          split_size_or_sections=batch_size // self.num_copies,
          dim=0,
      )
  
      # cross-correlation matrix
      correlation_matrix = torch.mm(
          embedding_normed_a.T, embedding_normed_b
      ) / (batch_size / self.num_copies)
  
      # Reduce cross-correlation matrices from all processes
      correlation_matrix = all_reduce_mean(correlation_matrix)
  
      # loss
      on_diag = (
          torch.diagonal(correlation_matrix)
          .add(-1)
          .pow(2)
          .sum()
          .mul(self.scale_loss)
      )
      off_diag = (
          self._off_diagonal(correlation_matrix)
          .pow(2)
          .sum()
          .mul(self.scale_loss)
      )
      loss = on_diag + self.lambda_ * off_diag
  
      return loss
  
file_path: vissl/vissl/losses/barlow_twins_loss.py
imports: import pprint
  import torch
  from classy_vision.generic.distributed_util import all_reduce_mean
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.utils.apex_helpers import SyncNormalizeFunction
  from vissl.utils.hydra_config import AttrDict
label_desc: Use image node_id 3 for calling the BarlowTwinsCriterion obj's underlying member method code with example usage: obj.forward(embedding) and returns: loss