class_docstr: Implementation of the loss described in 'Emerging Properties in
  Self-Supervised Vision Transformers'. [0]
  
  This implementation follows the code published by the authors. [1]
  It supports global and local image crops. A linear warmup schedule for the
  teacher temperature is implemented to stabilize training at the beginning.
  Centering is applied to the teacher output to avoid model collapse.
  
  - [0]: DINO, 2021, https://arxiv.org/abs/2104.14294
  - [1]: https://github.com/facebookresearch/dino
  
  Attributes:
      output_dim:
          Dimension of the model output.
      warmup_teacher_temp:
          Initial value of the teacher temperature. Should be decreased if the
          training loss does not decrease.
      teacher_temp:
          Final value of the teacher temperature after linear warmup. Values
          above 0.07 result in unstable behavior in most cases. Can be
          slightly increased to improve performance during finetuning.
      warmup_teacher_temp_epochs:
          Number of epochs for the teacher temperature warmup.
      student_temp:
          Temperature of the student.
      center_momentum:
          Momentum term for the center calculation.
  
  Examples:
  
      >>> # initialize loss function
      >>> loss_fn = DINOLoss(128)
      >>>
      >>> # generate a view of the images with a random transform
      >>> view = transform(images)
      >>>
      >>> # embed the view with a student and teacher model
      >>> teacher_out = teacher(view)
      >>> student_out = student(view)
      >>>
      >>> # calculate loss
      >>> loss = loss_fn([teacher_out], [student_out], epoch=0)
docstr: Cross-entropy between softmax outputs of the teacher and student
  networks.
  
  Args:
      teacher_out:
          List of view feature tensors from the teacher model. Each
          tensor is assumed to contain features from one view of the batch
          and have length batch_size.
      student_out:
          List of view feature tensors from the student model. Each tensor
          is assumed to contain features from one view of the batch and
          have length batch_size.
      epoch:
          The current training epoch.
  
  Returns:
      The average cross-entropy loss.
docstr_tok: ["Cross-entropy","between","softmax","outputs","of","the","teacher","and","student","networks",".","Args",":","teacher_out",":","List","of","view","feature","tensors","from","the","teacher","model",".","Each","tensor","is","assumed","to","contain","features","from","one","view","of","the","batch","and","have","length","batch_size",".","student_out",":","List","of","view","feature","tensors","from","the","student","model",".","Each","tensor","is","assumed","to","contain","features","from","one","view","of","the","batch","and","have","length","batch_size",".","epoch",":","The","current","training","epoch",".","Returns",":","The","average","cross-entropy","loss","."]
code: def forward(
      self,
      teacher_out: List[torch.Tensor],
      student_out: List[torch.Tensor],
      epoch: int,
  ) -> torch.Tensor:
      """Cross-entropy between softmax outputs of the teacher and student
      networks.
  
      Args:
          teacher_out:
              List of view feature tensors from the teacher model. Each
              tensor is assumed to contain features from one view of the batch
              and have length batch_size.
          student_out:
              List of view feature tensors from the student model. Each tensor
              is assumed to contain features from one view of the batch and
              have length batch_size.
          epoch:
              The current training epoch.
  
      Returns:
          The average cross-entropy loss.
  
      """
      # get teacher temperature
      if epoch < self.warmup_teacher_temp_epochs:
          teacher_temp = self.teacher_temp_schedule[epoch]
      else:
          teacher_temp = self.teacher_temp
  
      teacher_out = torch.stack(teacher_out)
      t_out = F.softmax(
          (teacher_out - self.center) / teacher_temp, dim=-1
      )
  
      student_out = torch.stack(student_out)
      s_out = F.log_softmax(student_out / self.student_temp, dim=-1)
  
      # calculate feature similarities where:
      # b = batch_size, t = n_views_teacher, s = n_views_student, d = output_dim
      # the diagonal is ignored as it contains features from the same views
      loss = -torch.einsum("tbd,sbd->ts", t_out, s_out)
      loss.fill_diagonal_(0)
  
      # number of loss terms, ignoring the diagonal
      n_terms = loss.numel() - loss.diagonal().numel()
      batch_size = teacher_out.shape[1]
      loss = loss.sum() / (n_terms * batch_size)
  
      self.update_center(teacher_out)
      return loss
  
file_path: lightly/lightly/loss/dino_loss.py
imports: from typing import List
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the DINOLoss obj's underlying member method code with example usage: obj.forward(teacher_out, student_out, epoch) and returns: loss