code: def test_subcenter_arcface_loss(self):
      batch_size = 64
      embedding_size = 32
      margin = 30
      scale = 64
      num_classes = 10
      for dtype in TEST_DTYPES:
          for sub_centers in [1, 3, 4]:
              loss_func = SubCenterArcFaceLoss(
                  margin=margin,
                  scale=scale,
                  num_classes=num_classes,
                  embedding_size=embedding_size,
                  sub_centers=sub_centers,
              )
  
              embeddings = (
                  torch.randn(batch_size, embedding_size)
                  .to(TEST_DEVICE)
                  .type(dtype)
              )
              labels = torch.randint(
                  low=0, high=num_classes, size=(batch_size,)
              ).to(TEST_DEVICE)
              # check if subcenters are included
              self.assertTrue(
                  loss_func.W.shape[1] == num_classes * sub_centers
              )
  
              loss = loss_func(embeddings, labels)
              loss.backward()
  
              weights = F.normalize(loss_func.W, p=2, dim=0)
              logits = torch.matmul(F.normalize(embeddings), weights)
              # include only closest sub centers
              logits = logits.view(-1, num_classes, sub_centers)
              logits, _ = logits.max(axis=2)
  
              for i, c in enumerate(labels):
                  acos = torch.acos(torch.clamp(logits[i, c], -1, 1))
                  logits[i, c] = torch.cos(
                      acos
                      + torch.tensor(
                          np.radians(margin), dtype=dtype
                      ).to(TEST_DEVICE)
                  )
  
              correct_loss = F.cross_entropy(
                  logits * scale, labels.to(TEST_DEVICE)
              )
  
              rtol = 1e-2 if dtype == torch.float16 else 1e-5
              self.assertTrue(
                  torch.isclose(loss, correct_loss, rtol=rtol)
              )
  
              if sub_centers == 1:
                  regular_arcface = ArcFaceLoss(
                      margin=margin,
                      scale=scale,
                      num_classes=num_classes,
                      embedding_size=embedding_size,
                  )
                  regular_arcface.W = loss_func.W
                  regular_arcface_loss = regular_arcface(
                      embeddings, labels
                  )
                  self.assertTrue(
                      torch.isclose(
                          loss, regular_arcface_loss, rtol=rtol
                      )
                  )
  
              # test get_logits
              logits_out = loss_func.get_logits(embeddings)
              self.assertTrue(
                  logits_out.shape
                  == torch.Size([batch_size, num_classes])
              )
              logits = torch.matmul(F.normalize(embeddings), weights)
              # include only closest sub centers
              logits = logits.view(-1, num_classes, sub_centers)
              logits_target, _ = logits.max(axis=2)
              self.assertTrue(
                  torch.allclose(logits_out, logits_target * scale)
              )
  
file_path: pytorch-metric-learning/tests/losses/test_subcenter_arcface_loss.py
imports: import math
  import unittest
  import numpy
  import torch
  import torch.nn.functional
  from pytorch_metric_learning.losses import ArcFaceLoss, SubCenterArcFaceLoss
  from ..None import TEST_DEVICE, TEST_DTYPES
label_desc: Use image node_id 1 for calling the TestSubCenterArcFaceLoss obj's underlying member method code with example usage: obj.test_subcenter_arcface_loss() without return types