class_docstr: Modified from Partial-FC
class_docstr_tok: ["Modified","from","Partial-FC"]
code: def __init__(
      self,
      rank,
      local_rank,
      world_size,
      batch_size,
      resume,
      margin_softmax,
      num_classes,
      sample_rate=1.0,
      embedding_size=512,
      prefix="./",
      cfg=None,
  ):
      super(VPL, self).__init__()
      #
      assert sample_rate == 1.0
      assert not resume
      self.num_classes: int = num_classes
      self.rank: int = rank
      self.local_rank: int = local_rank
      self.device: torch.device = torch.device(
          "cuda:{}".format(self.local_rank)
      )
      self.world_size: int = world_size
      self.batch_size: int = batch_size
      self.margin_softmax: callable = margin_softmax
      self.sample_rate: float = sample_rate
      self.embedding_size: int = embedding_size
      self.prefix: str = prefix
      self.num_local: int = num_classes // world_size + int(
          rank < num_classes % world_size
      )
      self.class_start: int = num_classes // world_size * rank + min(
          rank, num_classes % world_size
      )
      self.num_sample: int = int(self.sample_rate * self.num_local)
  
      self.weight_name = os.path.join(
          self.prefix, "rank_{}_softmax_weight.pt".format(self.rank)
      )
      self.weight_mom_name = os.path.join(
          self.prefix, "rank_{}_softmax_weight_mom.pt".format(self.rank)
      )
  
      self.weight = torch.normal(
          0,
          0.01,
          (self.num_local, self.embedding_size),
          device=self.device,
      )
      self.weight_mom: torch.Tensor = torch.zeros_like(self.weight)
      logging.info("softmax weight init successfully!")
      logging.info("softmax weight mom init successfully!")
      self.stream: torch.cuda.Stream = torch.cuda.Stream(local_rank)
  
      self.index = None
      self.update = lambda: 0
      self.sub_weight = Parameter(self.weight)
      self.sub_weight_mom = self.weight_mom
  
      # vpl variables
  
      self._iters = 0
      self.cfg = cfg
      self.vpl_mode = -1
      if self.cfg is not None:
          self.vpl_mode = self.cfg["mode"]
          if self.vpl_mode >= 0:
              self.register_buffer(
                  "queue",
                  torch.randn(
                      self.num_local,
                      self.embedding_size,
                      device=self.device,
                  ),
              )
              self.queue = normalize(self.queue)
              self.register_buffer(
                  "queue_iters",
                  torch.zeros(
                      (self.num_local,),
                      dtype=torch.long,
                      device=self.device,
                  ),
              )
              self.register_buffer(
                  "queue_lambda",
                  torch.zeros(
                      (self.num_local,),
                      dtype=torch.float32,
                      device=self.device,
                  ),
              )
  
file_path: insightface/recognition/vpl/vpl.py
imports: import logging
  import os
  import torch
  import torch.distributed
  from torch.nn import Module
  from torch.nn.functional import normalize, linear
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 1 to create a new VPL object from inherited base classes: Module with example: obj = VPL(rank, local_rank, world_size, batch_size, resume, margin_softmax, num_classes, sample_rate, embedding_size, prefix, cfg)