code: # 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  BaseDistillationMethod
  from solo.methods.base_for_AT_dual_bn import BaseDistillationATMethodDualBN
  from solo.utils.momentum import initialize_momentum_params
  from solo.utils.misc import gather
  
  from torchvision import models
  from solo.utils.metrics import accuracy_at_k, weighted_mean
  
  
  class MoCoV2KDATDualBN(BaseDistillationATMethodDualBN):
      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),
          )
  
          # Make StudentNetwork have same weight with Teacher
          self.projector.load_state_dict(self.projector_state_dict)
  
  
          # 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),
          )
          # initialize_momentum_params(self.projector, self.momentum_projector)
          self.momentum_projector.load_state_dict(self.projector_state_dict)
  
          # 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))
  
          # To remove both projector
          self.projector = nn.Identity()
          self.momentum_projector = nn.Identity()
  
      @staticmethod
      def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
          parent_parser = super(MoCoV2KDATDualBN, MoCoV2KDATDualBN).add_model_specific_args(parent_parser)
          parser = parent_parser.add_argument_group("mocov2_kd_at")
  
          # 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)
  
          # parser.add_argument("--limit_val_batches", type=float, default=0.2)
  
          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()}]
          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 backbone.
          """
  
          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 backbone and projector.
  
          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)
          z = F.normalize(self.projector(out["feats"]), dim=-1)
          return {**out, "z": z}
  
      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_large_crops
                  containing batches of images.
              batch_idx (int): index of the batch.
  
          Returns:
              torch.Tensor: total loss composed of MOCO loss and classification loss.
  
          """
  
  
          self.momentum_backbone.eval()
  
          # import ipdb; ipdb.set_trace()
          image_tau1, image_weak = batch[1]
          targets = batch[2]
  
          ############################################################################
          # Adversarial Training (DUAL BN)
          ############################################################################
  
          away_target = self.momentum_projector(self.momentum_backbone(image_weak))
          AE_generation_image = image_weak
          image_AE = self.generate_training_AE(AE_generation_image, away_target)
  
          student_logtis_clean = self.backbone(image_weak, "normal")
          student_logits_AE = self.backbone(image_AE, "pgd")
  
          adv_loss = -F.cosine_similarity(student_logits_AE, away_target).mean()
          adv_loss += -3*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
  
          ############################################################################
          # Adversarial Training (DUAL BN)
          ############################################################################
  
  
  
  
  
  
  
  
          ############################################################################
          # Adversarial Training (CAT)
          ############################################################################
  
          # away_target = self.momentum_projector(self.momentum_backbone(image_weak))
          # AE_generation_image = image_weak
          # image_AE = self.generate_training_AE(AE_generation_image, away_target)
  
          # image_CAT = torch.cat([image_weak, image_AE])
          # logits_all = self.backbone(image_CAT)
          # bs = image_weak.size(0)
          # student_logtis_clean = logits_all[:bs]
          # student_logits_AE = logits_all[bs:]
  
          # adv_loss = -F.cosine_similarity(student_logits_AE, away_target).mean()
          # adv_loss += -3*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
  
          ############################################################################
          # Adversarial Training (CAT)
          ############################################################################
  
  
  
  
  
  
          ############################################################################
          # Adversarial Training (CODE Experiment)
          ############################################################################
  
          # # NOTE: Check Wether to Make Student Start Same Weight with Teacher
          # # CODE Definition. 0: Weak augmentation 1: Tau1 augmentation
          # # CS_Teacher$1_Student-AE_CS_Teacher$2_Student$3, so case 1 is 110
          # # 1. CS_TeacherTau1_Student-AE_CS_TeacherTau1_StudentWeak
          # # 110
          # # with torch.no_grad():
          # #     logits_TeacherTau1 = self.projector(self.momentum_backbone(image_tau1))
  
          # # image_AE = self.generate_training_AE(image_weak, logits_TeacherTau1)
          # # student_logit = self.projector(self.backbone(image_AE))
          # # adv_loss = -F.cosine_similarity(student_logit, logits_TeacherTau1).mean()
  
          # # 2. CS_TeacherTau1_Student-AE_CS_TeacherTau1_StudentTau1
          # # 111
          # # with torch.no_grad():
          # #     logits_TeacherTau1 = self.projector(self.momentum_backbone(image_tau1))
  
          # # image_AE = self.generate_training_AE(image_tau1, logits_TeacherTau1)
          # # student_logit = self.projector(self.backbone(image_AE))
          # # adv_loss = -F.cosine_similarity(student_logit, logits_TeacherTau1).mean()
  
          # # CODE Based experiment
          # # 110
          # # 111
          # # experiment_code = [1, 0, 1]
          # # experiment_code = [1, 0, 0]
  
          # experiment_code = [0, 0, 0]
          # # experiment_code = [0, 1, 1]
          # # experiment_code = [0, 0, 1]
          # # experiment_code = [0, 1, 0]
  
          # with torch.no_grad():
          #     # away target
          #     if experiment_code[1] == 0:
          #         # away_target = self.projector(self.momentum_backbone(image_weak))
          #         away_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
          #     elif experiment_code[1] == 1:
          #         # away_target = self.projector(self.momentum_backbone(image_tau1))
          #         away_target = self.momentum_projector(self.momentum_backbone(image_tau1))
  
  
          #     # Loss learning target
          #     if experiment_code[0] == experiment_code[1]:
          #         learning_target = away_target
          #     elif experiment_code[0] == 0:
          #         # learning_target = self.projector(self.momentum_backbone(image_weak))
          #         learning_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
          #     elif experiment_code[0] == 1:
          #         # learning_target = self.projector(self.momentum_backbone(image_tau1))
          #         learning_target = self.momentum_projector(self.momentum_backbone(image_tau1))
  
  
          # if experiment_code[2] == 0:
          #     AE_generation_image = image_weak
          # elif experiment_code[2] == 1:
          #     AE_generation_image = image_tau1
  
          # image_AE = self.generate_training_AE(AE_generation_image, away_target)
          # student_logit = self.projector(self.backbone(image_AE))
          # adv_loss = -F.cosine_similarity(student_logit, learning_target).mean()
  
          ############################################################################
          # Adversarial Training (CODE Experiment)
          ############################################################################
  
  
  
  
  
          ############################################################################
          # Online clean classifier training
          ############################################################################
  
          # Bug Fix: train classifier using evaluation mode
          self.backbone.eval()
          outs_image_weak = self._base_shared_step(image_weak, targets, bn_name="pgd")
          self.backbone.train()
  
          metrics = {
              "train_class_loss": outs_image_weak["loss"],
              "train_acc1": outs_image_weak["acc1"],
              "train_acc5": outs_image_weak["acc5"],
          }
          class_loss_clean = outs_image_weak["loss"]
          self.log_dict(metrics, on_epoch=True)
          ############################################################################
          # Online clean classifier training
          ############################################################################
  
  
  
  
          ############################################################################
          # Online adv classifier training
          ############################################################################
  
          # Bug Fix: train classifier using evaluation mode
          # logits = self.classifier_adv(student_logits_AE)
  
          self.backbone.eval()
          AE_feats = self.backbone(image_AE, "pgd")
          logits = self.classifier_adv(AE_feats.detach())
          self.backbone.train()
          class_loss_adv = F.cross_entropy(logits, targets, ignore_index=-1)
  
          # handle when the number of classes is smaller than 5
          top_k_max = min(5, logits.size(1))
          acc1, acc5 = accuracy_at_k(logits, targets, top_k=(1, top_k_max))
  
          metrics = {
              "train_class_loss_adv_classifier": class_loss_adv,
              "train_acc1_adv_classifier": acc1,
              "train_acc5_adv_classifier": acc5,
          }
          class_loss_clean = outs_image_weak["loss"]
          self.log_dict(metrics, on_epoch=True)
          ############################################################################
          # Online adv classifier training
          ############################################################################
  
  
  
          self.log("adv_loss", adv_loss, on_epoch=True, sync_dist=True)
  
          return adv_loss + class_loss_adv + class_loss_clean
  
  
  
      def generate_training_AE(self, image: torch.Tensor, away_target: torch.Tensor):
          """
          images_org: weak aug
          away_target: from teacher
          """
  
          self.epsilon = 8/255.
          self.num_steps = 5
          self.step_size = 2/255.
  
          x_cl = image.clone().detach()
  
          # if self.rand:
          x_cl = x_cl + torch.zeros_like(image).uniform_(-self.epsilon, self.epsilon)
  
          # f_ori_proj = self.model(images_org).detach()
          # Change the attack process of model to eval
          self.backbone.eval()
  
          for i in range(self.num_steps):
              x_cl.requires_grad_()
              with torch.enable_grad():
                  f_proj = self.projector(self.backbone(x_cl, "pgd"))
  
                  # loss_contrast = -F.cosine_similarity(f_proj, f_ori_proj, dim=1).mean()
                  loss_contrast = -F.cosine_similarity(f_proj, away_target, dim=1).sum() *256
                  loss = loss_contrast
  
              # import ipdb ;ipdb.set_trace()
              grad_x_cl = torch.autograd.grad(loss, x_cl)[0]
              # grad_x_cl = torch.autograd.grad(loss, x_cl, grad_outputs=torch.ones_like(loss))[0]
              x_cl = x_cl.detach() + self.step_size * torch.sign(grad_x_cl.detach())
              x_cl = torch.min(torch.max(x_cl, image - self.epsilon), image + self.epsilon)
              x_cl = torch.clamp(x_cl, 0, 1)
  
          self.backbone.train()
  
          return x_cl
  
  
apis: torch.max torch.autograd.grad torch.nn.Linear torch.randn torch.nn.functional.normalize torch.zeros_like torch.no_grad torch.enable_grad torch.nn.functional.cross_entropy torch.nn.Identity torch.zeros torch.nn.ReLU torch.clamp torch.nn.functional.cosine_similarity