class_docstr: Distributed version of the NCE loss. It performs an "all_gather" to gather
  the allocated buffers like memory no a single gpu. For this, Pytorch distributed
  backend is used. If using NCCL, one must ensure that all the buffer are on GPU.
  This class supports training using both NCE and CrossEntropy (InfoNCE).
  
  This loss is used by NPID (https://arxiv.org/pdf/1805.01978.pdf), NPID++ and
  PIRL (https://arxiv.org/abs/1912.01991) approaches.
  
  Written by: Ishan Misra (imisra@fb.com)
  
  Config params:
      norm_embedding (bool):           whether to normalize embeddings
      temperature (float):             the temperature to apply to logits
      norm_constant (int):             Z parameter in the NCEAverage
      update_mem_with_emb_index (int): In case we have multiple embeddings used
                                       in the nce loss, specify which embedding
                                       to use to update the memory.
      loss_type (str):                 options are "nce" | "cross_entropy". Using the
                                       cross_entropy turns the loss into InfoNCE loss.
      loss_weights (List[float]):      if the NCE loss is computed between multiple pairs,
                                       we can set a loss weight per term can be used to weight
                                       different pair contributions differently
      negative_sampling_params:
          num_negatives (int):         how many negatives to contrast with
          type (str):                  how to select the negatives. options "random"
      memory_params:
          memory_size (int):           number of training samples as all the samples are
                                       stored in memory
          embedding_dim (int):         the projection head output dimension
          momentum (int):              momentum to use to update the memory
          norm_init (bool):            whether to L2 normalize the initialized memory bank
          update_mem_on_forward (bool): whether to update memory on the forward pass
      num_train_samples (int):         number of unique samples in the training dataset
class_docstr_tok: ["Distributed","version","of","the","NCE","loss",".","It","performs","an","``","all_gather","''","to","gather","the","allocated","buffers","like","memory","no","a","single","gpu",".","For","this",",","Pytorch","distributed","backend","is","used",".","If","using","NCCL",",","one","must","ensure","that","all","the","buffer","are","on","GPU",".","This","class","supports","training","using","both","NCE","and","CrossEntropy","(","InfoNCE",")",".","This","loss","is","used","by","NPID","(","https",":","\/\/arxiv.org\/pdf\/1805.01978.pdf",")",",","NPID++","and","PIRL","(","https",":","\/\/arxiv.org\/abs\/1912.01991",")","approaches",".","Written","by",":","Ishan","Misra","(","imisra","@","fb.com",")","Config","params",":","norm_embedding","(","bool",")",":","whether","to","normalize","embeddings","temperature","(","float",")",":","the","temperature","to","apply","to","logits","norm_constant","(","int",")",":","Z","parameter","in","the","NCEAverage","update_mem_with_emb_index","(","int",")",":","In","case","we","have","multiple","embeddings","used","in","the","nce","loss",",","specify","which","embedding","to","use","to","update","the","memory",".","loss_type","(","str",")",":","options","are","``","nce","''","|","``","cross_entropy","''",".","Using","the","cross_entropy","turns","the","loss","into","InfoNCE","loss",".","loss_weights","(","List","[","float","]",")",":","if","the","NCE","loss","is","computed","between","multiple","pairs",",","we","can","set","a","loss","weight","per","term","can","be","used","to","weight","different","pair","contributions","differently","negative_sampling_params",":","num_negatives","(","int",")",":","how","many","negatives","to","contrast","with","type","(","str",")",":","how","to","select","the","negatives",".","options","``","random","''","memory_params",":","memory_size","(","int",")",":","number","of","training","samples","as","all","the","samples","are","stored","in","memory","embedding_dim","(","int",")",":","the","projection","head","output","dimension","momentum","(","int",")",":","momentum","to","use","to","update","the","memory","norm_init","(","bool",")",":","whether","to","L2","normalize","the","initialized","memory","bank","update_mem_on_forward","(","bool",")",":","whether","to","update","memory","on","the","forward","pass","num_train_samples","(","int",")",":","number","of","unique","samples","in","the","training","dataset"]
code: def __init__(self, loss_config: AttrDict):
      super(NCELossWithMemory, self).__init__()
  
      self.loss_config = loss_config
      memory_params = self.loss_config.memory_params
      memory_params.memory_size = self.loss_config.num_train_samples
      assert is_pos_int(
          memory_params.memory_size
      ), f"Memory size must be positive: {memory_params.memory_size}"
  
      assert self.loss_config.loss_type in [
          "nce",
          "cross_entropy",
      ], f"Supported types are nce/cross_entropy. Found {self.loss_config.loss_type}"
  
      self.loss_type = self.loss_config.loss_type
  
      self.update_memory_on_forward = (
          memory_params.update_mem_on_forward
      )
      self.update_memory_emb_index = (
          self.loss_config.update_mem_with_emb_index
      )
      if self.update_memory_on_forward is False:
          # we have multiple embeddings used in NCE
          # but we update memory with only one of them
          assert self.update_memory_emb_index >= 0
  
      # first setup the NCEAverage method to get the scores of the output wrt
      # memory bank negatives
      self.nce_average = NCEAverage(
          memory_params=memory_params,
          negative_sampling_params=self.loss_config.negative_sampling_params,
          T=self.loss_config.temperature,
          Z=self.loss_config.norm_constant,
          loss_type=self.loss_type,
      )
  
      if self.loss_type == "nce":
          # setup the actual NCE loss
          self.nce_criterion = NCECriterion(
              self.loss_config.num_train_samples
          )
      elif self.loss_type == "cross_entropy":
          # cross-entropy loss. Also called InfoNCE
          self.xe_criterion = nn.CrossEntropyLoss()
  
      # other constants
      self.normalize_embedding = self.loss_config.norm_embedding
      self.loss_weights = self.loss_config.loss_weights
      self.init_sync_memory = False
      self.ignore_index = self.loss_config.get("ignore_index", -1)
  
imports: import logging
  import math
  import pprint
  from typing import List, Union
  import numpy
  import torch
  from classy_vision.generic.distributed_util import all_reduce_mean, gather_from_all, get_rank
  from classy_vision.generic.util import is_pos_int
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.config import AttrDict
label_desc: Use image node_id 1 to create a new NCELossWithMemory object from inherited base classes: ClassyLoss with example: obj = NCELossWithMemory(loss_config)