class_docstr: Loss used for DeepClusterV2 approach as provided in SwAV paper
  https://arxiv.org/abs/2006.09882
  
  Config params:
      DROP_LAST (bool): automatically inferred from DATA.TRAIN.DROP_LAST
      BATCHSIZE_PER_REPLICA (int): 256  # automatically inferred from
                                          DATA.TRAIN.BATCHSIZE_PER_REPLICA
      num_crops (int): 2                # automatically inferred from DATA.TRAIN.TRANSFORMS
      temperature (float): 0.1
      num_clusters (List[int]): [3000, 3000, 3000]
      kmeans_iters (int): 10
      crops_for_mb: [0]
      embedding_dim: 128
      num_train_samples (int): -1       # @auto-filled
class_docstr_tok: ["Loss","used","for","DeepClusterV2","approach","as","provided","in","SwAV","paper","https",":","\/\/arxiv.org\/abs\/2006.09882","Config","params",":","DROP_LAST","(","bool",")",":","automatically","inferred","from","DATA.TRAIN.DROP_LAST","BATCHSIZE_PER_REPLICA","(","int",")",":","256","#","automatically","inferred","from","DATA.TRAIN.BATCHSIZE_PER_REPLICA","num_crops","(","int",")",":","2","#","automatically","inferred","from","DATA.TRAIN.TRANSFORMS","temperature","(","float",")",":","0.1","num_clusters","(","List","[","int","]",")",":","[","3000",",","3000",",","3000","]","kmeans_iters","(","int",")",":","10","crops_for_mb",":","[","0","]","embedding_dim",":","128","num_train_samples","(","int",")",":","-1","#","@","auto-filled"]
code: def cluster_memory(self):
      self.start_idx = 0
      j = 0
      with torch.no_grad():
          for i_K, K in enumerate(self.num_clusters):
              # run distributed k-means
  
              # init centroids with elements from memory bank of rank 0
              centroids = torch.empty(K, self.embedding_dim).cuda(
                  non_blocking=True
              )
              if get_rank() == 0:
                  random_idx = torch.randperm(
                      len(self.local_memory_embeddings[j])
                  )[:K]
                  assert (
                      len(random_idx) >= K
                  ), "please reduce the number of centroids"
                  centroids = self.local_memory_embeddings[j][
                      random_idx
                  ]
              dist.broadcast(centroids, 0)
  
              for n_iter in range(self.nmb_kmeans_iters + 1):
                  # E step
                  dot_products = torch.mm(
                      self.local_memory_embeddings[j], centroids.t()
                  )
                  _, assignments = dot_products.max(dim=1)
  
                  # finish
                  if n_iter == self.nmb_kmeans_iters:
                      break
  
                  # M step
                  where_helper = get_indices_sparse(
                      assignments.cpu().numpy()
                  )
                  counts = torch.zeros(K).cuda(non_blocking=True).int()
                  emb_sums = torch.zeros(K, self.embedding_dim).cuda(
                      non_blocking=True
                  )
                  for k in range(len(where_helper)):
                      if len(where_helper[k][0]) > 0:
                          emb_sums[k] = torch.sum(
                              self.local_memory_embeddings[j][
                                  where_helper[k][0]
                              ],
                              dim=0,
                          )
                          counts[k] = len(where_helper[k][0])
                  all_reduce_sum(counts)
                  mask = counts > 0
                  all_reduce_sum(emb_sums)
                  centroids[mask] = emb_sums[mask] / counts[
                      mask
                  ].unsqueeze(1)
  
                  # normalize centroids
                  centroids = nn.functional.normalize(
                      centroids, dim=1, p=2
                  )
  
              getattr(self, "centroids" + str(i_K)).copy_(centroids)
              # gather the assignments
              assignments_all = gather_from_all(assignments)
              indexes_all = gather_from_all(self.local_memory_index)
              self.assignments[i_K] = -100
              self.assignments[i_K][indexes_all] = assignments_all
  
              j = (j + 1) % self.nmb_mbs
  
      logging.info(
          f"Rank: {get_rank()}, clustering of the memory bank done"
      )
  
file_path: vissl/vissl/losses/deepclusterv2_loss.py
imports: import logging
  import math
  import pprint
  import torch
  import torch.distributed
  from classy_vision.generic.distributed_util import all_reduce_sum, gather_from_all, get_rank, get_world_size
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.config import AttrDict
  from vissl.utils.misc import get_indices_sparse
label_desc: Use image node_id 6 for calling the DeepClusterV2Loss obj's underlying member method code with example usage: obj.cluster_memory() without return types