class_docstr: https://arxiv.org/abs/2203.15565
  A distributed sparsely updating variant of the FC layer, named Partial FC (PFC).
  When sample rate less than 1, in each iteration, positive class centers and a random subset of
  negative class centers are selected to compute the margin-based softmax loss, all class
  centers are still maintained throughout the whole training process, but only a subset is
  selected and updated in each iteration.
  .. note::
      When sample rate equal to 1, Partial FC is equal to model parallelism(default sample rate is 1).
  Example:
  --------
  >>> module_pfc = PartialFC(embedding_size=512, num_classes=8000000, sample_rate=0.2)
  >>> for img, labels in data_loader:
  >>>     embeddings = net(img)
  >>>     loss = module_pfc(embeddings, labels)
  >>>     loss.backward()
  >>>     optimizer.step()
class_docstr_tok: ["https",":","\/\/arxiv.org\/abs\/2203.15565","A","distributed","sparsely","updating","variant","of","the","FC","layer",",","named","Partial","FC","(","PFC",")",".","When","sample","rate","less","than","1",",","in","each","iteration",",","positive","class","centers","and","a","random","subset","of","negative","class","centers","are","selected","to","compute","the","margin-based","softmax","loss",",","all","class","centers","are","still","maintained","throughout","the","whole","training","process",",","but","only","a","subset","is","selected","and","updated","in","each","iteration",".","..","note",":",":","When","sample","rate","equal","to","1",",","Partial","FC","is","equal","to","model","parallelism","(","default","sample","rate","is","1",")",".","Example",":","--","--","--","--",">",">",">","module_pfc","=","PartialFC","(","embedding_size=512",",","num_classes=8000000",",","sample_rate=0.2",")",">",">",">","for","img",",","labels","in","data_loader",":",">",">",">","embeddings","=","net","(","img",")",">",">",">","loss","=","module_pfc","(","embeddings",",","labels",")",">",">",">","loss.backward","(",")",">",">",">","optimizer.step","(",")"]
docstr: Parameters:
  ----------
  local_embeddings: torch.Tensor
      feature embeddings on each GPU(Rank).
  local_labels: torch.Tensor
      labels on each GPU(Rank).
  Returns:
  -------
  loss: torch.Tensor
      pass
docstr_tok: ["Parameters",":","--","--","--","--","--","local_embeddings",":","torch.Tensor","feature","embeddings","on","each","GPU","(","Rank",")",".","local_labels",":","torch.Tensor","labels","on","each","GPU","(","Rank",")",".","Returns",":","--","--","--","-","loss",":","torch.Tensor","pass"]
code: def forward(
      self,
      local_embeddings: torch.Tensor,
      local_labels: torch.Tensor,
  ):
      """
      Parameters:
      ----------
      local_embeddings: torch.Tensor
          feature embeddings on each GPU(Rank).
      local_labels: torch.Tensor
          labels on each GPU(Rank).
      Returns:
      -------
      loss: torch.Tensor
          pass
      """
      local_labels.squeeze_()
      local_labels = local_labels.long()
  
      batch_size = local_embeddings.size(0)
      if self.last_batch_size == 0:
          self.last_batch_size = batch_size
      assert (
          self.last_batch_size == batch_size
      ), f"last batch size do not equal current batch size: {self.last_batch_size} vs {batch_size}"
  
      _gather_embeddings = [
          torch.zeros((batch_size, self.embedding_size)).cuda()
          for _ in range(self.world_size)
      ]
      _gather_labels = [
          torch.zeros(batch_size).long().cuda()
          for _ in range(self.world_size)
      ]
      _list_embeddings = AllGather(
          local_embeddings, *_gather_embeddings
      )
      distributed.all_gather(_gather_labels, local_labels)
  
      embeddings = torch.cat(_list_embeddings)
      labels = torch.cat(_gather_labels)
  
      labels = labels.view(-1, 1)
      index_positive = (self.class_start <= labels) & (
          labels < self.class_start + self.num_local
      )
      labels[~index_positive] = -1
      labels[index_positive] -= self.class_start
  
      if self.sample_rate < 1:
          weight = self.sample(labels, index_positive)
      else:
          weight = self.weight
  
      with torch.cuda.amp.autocast(self.fp16):
          norm_embeddings = normalize(embeddings)
          norm_weight_activated = normalize(weight)
          logits = linear(norm_embeddings, norm_weight_activated)
      if self.fp16:
          logits = logits.float()
      logits = logits.clamp(-1, 1)
  
      logits = self.margin_softmax(logits, labels)
      loss = self.dist_cross_entropy(logits, labels)
      return loss
  
imports: import math
  from typing import Callable
  import torch
  from torch import distributed
  from torch.nn.functional import linear, normalize
label_desc: Use image node_id 3 for calling the PartialFC_V2 obj's underlying member method code with example usage: obj.forward(local_embeddings, local_labels) and returns: loss