text: # Copyright 2021 solo-learn development team.
  
  # Permission is hereby granted, free of charge, to any person obtaining a copy of
  # this software and associated documentation files (the "Software"), to deal in
  # the Software without restriction, including without limitation the rights to use,
  # copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
  # Software, and to permit persons to whom the Software is furnished to do so,
  # subject to the following conditions:
  
  # The above copyright notice and this permission notice shall be included in all copies
  # or substantial portions of the Software.
  
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
  # INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
  # PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
  # FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
  # OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  # DEALINGS IN THE SOFTWARE.
  
  import argparse
  from typing import Any, Dict, List, Sequence, Tuple
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from solo.losses.moco import moco_loss_func
  from solo.methods.base import BaseMomentumMethod
  from solo.utils.momentum import initialize_momentum_params
  from solo.utils.misc import gather
  
  def rotate_tensors(X: torch.Tensor, rot: torch.Tensor) -> torch.Tensor:
      mask = rot == 1
      X[mask] = torch.flip(X[mask], dims=(3,)).transpose(dim0=2, dim1=3)
      mask = rot == 2
      X[mask] = torch.flip(X[mask], dims=(2,)).flip(dims=(3,))
      mask = rot == 3
      X[mask] = torch.flip(X[mask], dims=(2,)).transpose(dim0=2, dim1=3)
      return X
  
  # 两个相互独立的旋转标签 rev1
  # def rotate_batch(batch):
  #     _, X, cls_label = batch
  #     rot_label_0 = torch.randint(0, 4, cls_label.shape, device=cls_label.device)
  #     rot_label_1 = torch.randint(0, 4, cls_label.shape, device=cls_label.device)
  #     X[0] = rotate_tensors(X[0], rot_label_0)
  #     X[1] = rotate_tensors(X[1], rot_label_1)
  #     batch = (_, X, cls_label)
  #     return batch
  
  # 同样的旋转标签 rev2
  def rotate_batch(batch):
      _, X, cls_label = batch
      rot_label_0 = torch.randint(0, 4, cls_label.shape, device=cls_label.device)
      X[0] = rotate_tensors(X[0], rot_label_0)
      X[1] = rotate_tensors(X[1], rot_label_0)
      batch = (_, X, cls_label)
      return batch
  
  
  class AR_Rotation_MoCoV2Plus(BaseMomentumMethod):
      queue: torch.Tensor
  
      def __init__(
          self,
          proj_output_dim: int,
          proj_hidden_dim: int,
          temperature: float,
          queue_size: int,
          **kwargs
      ):
          """Implements MoCo V2+ (https://arxiv.org/abs/2011.10566).
  
          Args:
              proj_output_dim (int): number of dimensions of projected features.
              proj_hidden_dim (int): number of neurons of the hidden layers of the projector.
              temperature (float): temperature for the softmax in the contrastive loss.
              queue_size (int): number of samples to keep in the queue.
          """
  
          super().__init__(**kwargs)
  
          self.temperature = temperature
          self.queue_size = queue_size
  
          # projector
          self.projector = nn.Sequential(
              nn.Linear(self.features_dim, proj_hidden_dim),
              nn.ReLU(),
              nn.Linear(proj_hidden_dim, proj_output_dim),
          )
  
          # momentum projector
          self.momentum_projector = nn.Sequential(
              nn.Linear(self.features_dim, proj_hidden_dim),
              nn.ReLU(),
              nn.Linear(proj_hidden_dim, proj_output_dim),
          )
          non_linear = False
          if non_linear:
              # a_rot non-linear classifier
              self.a_rot_classifier = nn.Sequential(
                  nn.Linear(self.features_dim, proj_hidden_dim),
                  nn.ReLU(),
                  nn.Linear(proj_hidden_dim, 4)
              )
          else:
              # a_rot linear classifier
              self.a_rot_classifier = nn.Linear(self.features_dim, 4)
  
          initialize_momentum_params(self.projector, self.momentum_projector)
  
          # create the queue
          self.register_buffer("queue", torch.randn(2, proj_output_dim, queue_size))
          self.queue = nn.functional.normalize(self.queue, dim=1)
          self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
  
      @staticmethod
      def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
          parent_parser = super(AR_Rotation_MoCoV2Plus, AR_Rotation_MoCoV2Plus).add_model_specific_args(parent_parser)
          parser = parent_parser.add_argument_group("mocov2plus")
  
          # projector
          parser.add_argument("--proj_output_dim", type=int, default=128)
          parser.add_argument("--proj_hidden_dim", type=int, default=2048)
  
          # parameters
          parser.add_argument("--temperature", type=float, default=0.1)
  
          # queue settings
          parser.add_argument("--queue_size", default=65536, type=int)
  
          return parent_parser
  
      @property
      def learnable_params(self) -> List[dict]:
          """Adds projector parameters together with parent's learnable parameters.
  
          Returns:
              List[dict]: list of learnable parameters.
          """
  
          extra_learnable_params = [
              {"params": self.projector.parameters()},
              {"params": self.a_rot_classifier.parameters()}
              ]
          return super().learnable_params + extra_learnable_params
  
      @property
      def momentum_pairs(self) -> List[Tuple[Any, Any]]:
          """Adds (projector, momentum_projector) to the parent's momentum pairs.
  
          Returns:
              List[Tuple[Any, Any]]: list of momentum pairs.
          """
  
          extra_momentum_pairs = [(self.projector, self.momentum_projector)]
          return super().momentum_pairs + extra_momentum_pairs
  
      @torch.no_grad()
      def _dequeue_and_enqueue(self, keys: torch.Tensor):
          """Adds new samples and removes old samples from the queue in a fifo manner.
  
          Args:
              keys (torch.Tensor): output features of the momentum encoder.
          """
  
          batch_size = keys.shape[1]
          ptr = int(self.queue_ptr)  # type: ignore
          assert self.queue_size % batch_size == 0  # for simplicity
  
          # replace the keys at ptr (dequeue and enqueue)
          keys = keys.permute(0, 2, 1)
          self.queue[:, :, ptr : ptr + batch_size] = keys
          ptr = (ptr + batch_size) % self.queue_size  # move pointer
          self.queue_ptr[0] = ptr  # type: ignore
  
      def forward(self, X: torch.Tensor, *args, **kwargs) -> Dict[str, Any]:
          """Performs the forward pass of the online encoder and the online projection.
  
          Args:
              X (torch.Tensor): a batch of images in the tensor format.
  
          Returns:
              Dict[str, Any]: a dict containing the outputs of the parent and the projected features.
          """
  
          out = super().forward(X, *args, **kwargs)
          q = F.normalize(self.projector(out["feats"]), dim=-1)
          return {**out, "q": q}
  
      def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
          """
          Training step for MoCo reusing BaseMomentumMethod training step.
  
          Args:
              batch (Sequence[Any]): a batch of data in the
                  format of [img_indexes, [X], Y], where [X] is a list of size self.num_crops
                  containing batches of images.
              batch_idx (int): index of the batch.
  
          Returns:
              torch.Tensor: total loss composed of MOCO loss and classification loss.
  
          """
          # 将输入数据做旋转
          batch = rotate_batch(batch)
          out = super().training_step(batch, batch_idx)
          class_loss = out["loss"]
          feats1, feats2 = out["feats"]
          momentum_feats1, momentum_feats2 = out["momentum_feats"]
  
          q1 = self.projector(feats1)
          q2 = self.projector(feats2)
          q1 = F.normalize(q1, dim=-1)
          q2 = F.normalize(q2, dim=-1)
  
          with torch.no_grad():
              k1 = self.momentum_projector(momentum_feats1)
              k2 = self.momentum_projector(momentum_feats2)
              k1 = F.normalize(k1, dim=-1)
              k2 = F.normalize(k2, dim=-1)
  
          # ------- contrastive loss -------
          # symmetric
          queue = self.queue.clone().detach()
          nce_loss = (
              moco_loss_func(q1, k2, queue[1], self.temperature)
              + moco_loss_func(q2, k1, queue[0], self.temperature)
          ) / 2
  
          # ------- update queue -------
          keys = torch.stack((gather(k1), gather(k2)))
          self._dequeue_and_enqueue(keys)
  
          self.log("train_nce_loss", nce_loss, on_epoch=True, sync_dist=True)
  
          return nce_loss + class_loss
  