docstr: Test if freeze last layer cancels backprop.
docstr_tok: ["Test","if","freeze","last","layer","cancels","backprop","."]
code: def test_dino_projection_head_freeze_last_layer(self, seed=0):
      """Test if freeze last layer cancels backprop."""
      torch.manual_seed(seed)
      for norm_last_layer in [False, True]:
          for freeze_last_layer in range(-1, 3):
              head = DINOProjectionHead(
                  input_dim=4,
                  hidden_dim=4,
                  output_dim=4,
                  bottleneck_dim=4,
                  freeze_last_layer=freeze_last_layer,
                  norm_last_layer=norm_last_layer,
              )
              optimizer = torch.optim.SGD(head.parameters(), lr=1)
              criterion = lightly.loss.DINOLoss(output_dim=4)
              # Store initial weights of last layer
              initial_data = [
                  param.data.detach().clone()
                  for param in head.last_layer.parameters()
              ]
              for epoch in range(5):
                  with self.subTest(
                      f"norm_last_layer={norm_last_layer}, "
                      f"freeze_last_layer={freeze_last_layer}, "
                      f"epoch={epoch}"
                  ):
                      views = [torch.rand((3, 4)) for _ in range(2)]
                      teacher_out = [head(view) for view in views]
                      student_out = [head(view) for view in views]
                      loss = criterion(
                          teacher_out, student_out, epoch=epoch
                      )
                      optimizer.zero_grad()
                      loss.backward()
                      head.cancel_last_layer_gradients(
                          current_epoch=epoch
                      )
                      optimizer.step()
                      params = head.last_layer.parameters()
                      # Verify that weights have (not) changed depending on epoch.
                      for param, init_data in zip(params, initial_data):
                          if param.requires_grad:
                              are_same = torch.allclose(
                                  param.data, init_data
                              )
                              if epoch >= freeze_last_layer:
                                  self.assertFalse(are_same)
                              else:
                                  self.assertTrue(are_same)
  
file_path: lightly/tests/models/test_ProjectionHeads.py
imports: import unittest
  import torch
  import lightly
  from lightly.models.modules.heads import BarlowTwinsProjectionHead, BYOLPredictionHead, BYOLProjectionHead, DINOProjectionHead, MoCoProjectionHead, MSNProjectionHead, NNCLRPredictionHead, NNCLRProjectionHead, SimCLRProjectionHead, SimSiamPredictionHead, SimSiamProjectionHead, SwaVProjectionHead, SwaVPrototypes, TiCoProjectionHead, VicRegLLocalProjectionHead
label_desc: Use image node_id 10 for calling the TestProjectionHeads obj's underlying member method code with example usage: obj.test_dino_projection_head_freeze_last_layer(seed) without return types