class_docstr: This is the loss which was proposed in the "Momentum Contrast
  for Unsupervised Visual Representation Learning" paper, from Kaiming He et al.
  See http://arxiv.org/abs/1911.05722 for details
  and https://github.com/facebookresearch/moco for a reference implementation, reused here
  
  Config params:
      embedding_dim (int): head output output dimension
      queue_size (int): number of elements in queue
      momentum (float): encoder momentum value for the update
      temperature (float): temperature to use on the logits
class_docstr_tok: ["This","is","the","loss","which","was","proposed","in","the","``","Momentum","Contrast","for","Unsupervised","Visual","Representation","Learning","''","paper",",","from","Kaiming","He","et","al",".","See","http",":","\/\/arxiv.org\/abs\/1911.05722","for","details","and","https",":","\/\/github.com\/facebookresearch\/moco","for","a","reference","implementation",",","reused","here","Config","params",":","embedding_dim","(","int",")",":","head","output","output","dimension","queue_size","(","int",")",":","number","of","elements","in","queue","momentum","(","float",")",":","encoder","momentum","value","for","the","update","temperature","(","float",")",":","temperature","to","use","on","the","logits"]
docstr: Given the encoder queries, the key and the queue of the previous queries,
  compute the cross entropy loss for this batch
  
  Args:
      query: output of the encoder given the current batch
  
  Returns:
      loss
docstr_tok: ["Given","the","encoder","queries",",","the","key","and","the","queue","of","the","previous","queries",",","compute","the","cross","entropy","loss","for","this","batch","Args",":","query",":","output","of","the","encoder","given","the","current","batch","Returns",":","loss"]
code: def forward(
      self, query: torch.Tensor, *args, **kwargs
  ) -> torch.Tensor:
      """
      Given the encoder queries, the key and the queue of the previous queries,
      compute the cross entropy loss for this batch
  
      Args:
          query: output of the encoder given the current batch
  
      Returns:
          loss
      """
  
      if not self.initialized:
          self.queue = self.queue.to(query.device)
          self.initialized = True
  
      # --
      # Normalize the encoder raw outputs
      query = nn.functional.normalize(query, dim=1)
  
      # --
      # Compute all the logits and the expected labels
      # Einstein sum is used in MoCo, deemed more intuitive.
      # Another option is `torch.diag(torch.matmul(query, self.key.T))`
  
      # positive logits: Nx1
      l_pos = torch.einsum("nc,nc->n", [query, self.key]).unsqueeze(-1)
  
      # negative logits: NxK
      l_neg = torch.einsum(
          "nc,ck->nk", [query, self.queue.clone().detach()]
      )
  
      # logits: Nx(1+K)
      logits = torch.cat([l_pos, l_neg], dim=1)
  
      # apply temperature
      logits /= self.loss_config.temperature
  
      # labels: positives are the first rank.
      # This is essentially a classification problem alongside all the samples
      labels = torch.zeros(logits.shape[0], dtype=torch.long).to(
          query.device
      )
  
      # ---
      # Update the queue for the next time
      self._dequeue_and_enqueue(self.key)
  
      # ---
      # Then just apply the XELoss
      return self.criterion(logits, labels)
  
file_path: vissl/vissl/losses/moco_loss.py
imports: import logging
  import pprint
  from collections import namedtuple
  import torch
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.utils.misc import concat_all_gather
label_desc: Use image node_id 4 for calling the MoCoLoss obj's underlying member method code with example usage: obj.forward(query) and returns: self