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: Paramenters:
  -----------
  embedding_size: int
      The dimension of embedding, required
  num_classes: int
      Total number of classes, required
  sample_rate: float
      The rate of negative centers participating in the calculation, default is 1.0.
docstr_tok: ["Paramenters",":","--","--","--","--","--","-","embedding_size",":","int","The","dimension","of","embedding",",","required","num_classes",":","int","Total","number","of","classes",",","required","sample_rate",":","float","The","rate","of","negative","centers","participating","in","the","calculation",",","default","is","1.0","."]
code: def __init__(
      self,
      margin_loss: Callable,
      embedding_size: int,
      num_classes: int,
      sample_rate: float = 1.0,
      fp16: bool = False,
  ):
      """
      Paramenters:
      -----------
      embedding_size: int
          The dimension of embedding, required
      num_classes: int
          Total number of classes, required
      sample_rate: float
          The rate of negative centers participating in the calculation, default is 1.0.
      """
      super(PartialFC_V2, self).__init__()
      assert (
          distributed.is_initialized()
      ), "must initialize distributed before create this"
      self.rank = distributed.get_rank()
      self.world_size = distributed.get_world_size()
  
      self.dist_cross_entropy = DistCrossEntropy()
      self.embedding_size = embedding_size
      self.sample_rate: float = sample_rate
      self.fp16 = fp16
      self.num_local: int = num_classes // self.world_size + int(
          self.rank < num_classes % self.world_size
      )
      self.class_start: int = (
          num_classes // self.world_size * self.rank
          + min(self.rank, num_classes % self.world_size)
      )
      self.num_sample: int = int(self.sample_rate * self.num_local)
      self.last_batch_size: int = 0
  
      self.is_updated: bool = True
      self.init_weight_update: bool = True
      self.weight = torch.nn.Parameter(
          torch.normal(0, 0.01, (self.num_local, embedding_size))
      )
  
      # margin_loss
      if isinstance(margin_loss, Callable):
          self.margin_softmax = margin_loss
      else:
          raise
  
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 1 to create a new PartialFC_V2 object from inherited base classes: torch.nn with example: obj = PartialFC_V2(margin_loss, embedding_size, num_classes, sample_rate, fp16)