code: def test_dino_loss_equal_to_original(self):
      def test(
          batch_size=3,
          n_global=2,  # number of global views
          n_local=6,  # number of local views
          output_dim=4,
          warmup_teacher_temp=0.04,
          teacher_temp=0.04,
          warmup_teacher_temp_epochs=30,
          student_temp=0.1,
          center_momentum=0.9,
          epoch=0,
          n_epochs=100,
      ):
          """Runs test with the given input parameters."""
          with self.subTest(
              f"batch_size={batch_size}, n_global={n_global}, "
              f"n_local={n_local}, output_dim={output_dim}, "
              f"warmup_teacher_temp={warmup_teacher_temp}, "
              f"teacher_temp={teacher_temp}, "
              f"warmup_teacher_temp_epochs={warmup_teacher_temp_epochs}, "
              f"student_temp={student_temp}, "
              f"center_momentum={center_momentum}, epoch={epoch}, "
              f"n_epochs={n_epochs}"
          ):
              loss_fn = DINOLoss(
                  output_dim=output_dim,
                  warmup_teacher_temp=warmup_teacher_temp,
                  teacher_temp=teacher_temp,
                  warmup_teacher_temp_epochs=warmup_teacher_temp_epochs,
                  student_temp=student_temp,
                  center_momentum=center_momentum,
              )
  
              orig_loss_fn = OriginalDINOLoss(
                  out_dim=output_dim,
                  ncrops=n_global + n_local,
                  teacher_temp=teacher_temp,
                  warmup_teacher_temp=warmup_teacher_temp,
                  warmup_teacher_temp_epochs=warmup_teacher_temp_epochs,
                  nepochs=n_epochs,
                  student_temp=student_temp,
                  center_momentum=center_momentum,
              )
  
              # Create dummy single layer network. We use this to verify
              # that the gradient backprop works properly.
              teacher = torch.nn.Linear(output_dim, output_dim)
              deactivate_requires_grad(teacher)
              student = torch.nn.Linear(output_dim, output_dim)
              orig_teacher = copy.deepcopy(teacher)
              orig_student = copy.deepcopy(student)
  
              optimizer = torch.optim.SGD(student.parameters(), lr=1)
              orig_optimizer = torch.optim.SGD(
                  orig_student.parameters(), lr=1
              )
  
              # Create fake output
              teacher_out = self.generate_output(
                  batch_size=batch_size,
                  n_views=n_global,
                  output_dim=output_dim,
                  seed=0,
              )
              student_out = self.generate_output(
                  batch_size=batch_size,
                  n_views=n_global + n_local,
                  output_dim=output_dim,
                  seed=1,
              )
  
              # Clone input tensors
              orig_teacher_out = torch.cat(teacher_out)
              orig_teacher_out = orig_teacher_out.detach().clone()
              orig_student_out = torch.cat(student_out)
              orig_student_out = orig_student_out.detach().clone()
  
              # Forward pass
              teacher_out = [teacher(view) for view in teacher_out]
              student_out = [student(view) for view in student_out]
              orig_teacher_out = orig_teacher(orig_teacher_out)
              orig_student_out = orig_student(orig_student_out)
  
              # Calculate loss
              loss = loss_fn(
                  teacher_out=teacher_out,
                  student_out=student_out,
                  epoch=epoch,
              )
              orig_loss = orig_loss_fn(
                  student_output=orig_student_out,
                  teacher_output=orig_teacher_out,
                  epoch=epoch,
              )
  
              # Backward pass and optimizer step
              optimizer.zero_grad()
              orig_optimizer.zero_grad()
              loss.backward()
              orig_loss.backward()
              optimizer.step()
              orig_optimizer.step()
  
              # Loss and loss center should be equal
              center = loss_fn.center.squeeze()
              orig_center = orig_loss_fn.center.squeeze()
              self.assertTrue(torch.allclose(center, orig_center))
              self.assertTrue(torch.allclose(loss, orig_loss))
  
              # Parameters of network should be equal after backward pass
              for param, orig_param in zip(
                  student.parameters(), orig_student.parameters()
              ):
                  self.assertTrue(torch.allclose(param, orig_param))
              for param, orig_param in zip(
                  teacher.parameters(), orig_teacher.parameters()
              ):
                  self.assertTrue(torch.allclose(param, orig_param))
  
      def test_all(**kwargs):
          """Tests all combinations of the input parameters"""
          parameters = []
          for name, values in kwargs.items():
              parameters.append([(name, value) for value in values])
          # parameters = [
          #   [(param1, val11), (param1, val12), ..],
          #   [(param2, val21), (param2, val22), ..],
          #   ...
          # ]
  
          for params in itertools.product(*parameters):
              # params = [(param1, value1), (param2, value2), ...]
              test(**dict(params))
  
      # test input sizes
      test_all(
          batch_size=np.arange(1, 4),
          n_local=np.arange(0, 4),
          output_dim=np.arange(1, 4),
      )
      # test teacher temp warmup
      test_all(
          warmup_teacher_temp=[0.01, 0.04, 0.07],
          teacher_temp=[0.01, 0.04, 0.07],
          warmup_teacher_temp_epochs=[0, 1, 10],
          epoch=[0, 1, 10, 20],
      )
      # test other params
      test_all(
          student_temp=[0.05, 0.1, 0.2],
          center_momentum=[0.5, 0.9, 0.95],
      )
  
file_path: lightly/tests/loss/test_DINOLoss.py
imports: import copy
  import itertools
  import unittest
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn
  from lightly.loss import DINOLoss
  from lightly.models.utils import deactivate_requires_grad
label_desc: Use image node_id 2 for calling the TestDINOLoss obj's underlying member method code with example usage: obj.test_dino_loss_equal_to_original() without return types