docstr: Parameters
  ----------
  num_sample: int
      The number of sampled class center.
  num_local: int
      The number of class center storage in this rank(CPU/GPU).
  rank: int
      Unique process(GPU) ID from 0 to size - 1.
  local_rank: int
      Unique process(GPU) ID within the server from 0 to 7.
  embedding_size: int
      The feature dimension.
  prefix_dir: str
      Path prefix of model dir.
  gpu: bool
      If True, class center and class center mom will storage in GPU.
code: def __init__(
      self,
      num_sample,
      num_local,
      rank,
      local_rank,
      embedding_size,
      prefix,
      gpu=True,
  ):
      """
      Parameters
      ----------
      num_sample: int
          The number of sampled class center.
      num_local: int
          The number of class center storage in this rank(CPU/GPU).
      rank: int
          Unique process(GPU) ID from 0 to size - 1.
      local_rank: int
          Unique process(GPU) ID within the server from 0 to 7.
      embedding_size: int
          The feature dimension.
      prefix_dir: str
          Path prefix of model dir.
      gpu: bool
          If True, class center and class center mom will storage in GPU.
      """
      self.num_sample = num_sample
      self.num_local = num_local
      self.rank = rank
      self.embedding_size = embedding_size
      self.gpu = gpu
      self.prefix = prefix
  
      if self.gpu:
          context = mx.gpu(local_rank)
      else:
          context = mx.cpu()
  
      # In order to apply update, weight and momentum should be storage.
      self.weight = nd.random_normal(
          loc=0,
          scale=0.01,
          shape=(self.num_local, self.embedding_size),
          ctx=context,
      )
      self.weight_mom = nd.zeros_like(self.weight)
  
      # Sampler object
      self.weight_index_sampler = WeightIndexSampler(
          num_sample, num_local, rank
      )
  
file_path: insightface/recognition/partial_fc/mxnet/memory_bank.py
imports: import os
  import numpy
  from mxnet import nd
  import mxnet
  from memory_samplers import WeightIndexSampler
label_desc: Use image node_id 1 to create a new MemoryBank object from inherited base classes: object with example: obj = MemoryBank(num_sample, num_local, rank, local_rank, embedding_size, prefix, gpu)