class_docstr: Class performing adversarial training following Oracle Aligned Adversarial Training (OAAT) protocol.
  
  | Paper link: https://link.springer.com/chapter/10.1007/978-3-031-20065-6_18
class_docstr_tok: ["Class","performing","adversarial","training","following","Oracle","Aligned","Adversarial","Training","(","OAAT",")","protocol",".","|","Paper","link",":","https",":","\/\/link.springer.com\/chapter\/10.1007\/978-3-031-20065-6_18"]
docstr: Perform the operations of OAAT for a batch of data.
  See class documentation for more information on the exact procedure.
  
  :param i_epoch: training epoch number.
  :param nb_epochs: total training epochs.
  :param batch_id: batch_id of input data.
  :param x_batch: batch of x.
  :param y_batch: batch of y.
  :return: tuple containing batch data loss, batch data accuracy and number of samples in the batch
docstr_tok: ["Perform","the","operations","of","OAAT","for","a","batch","of","data",".","See","class","documentation","for","more","information","on","the","exact","procedure",".",":","param","i_epoch",":","training","epoch","number",".",":","param","nb_epochs",":","total","training","epochs",".",":","param","batch_id",":","batch_id","of","input","data",".",":","param","x_batch",":","batch","of","x.",":","param","y_batch",":","batch","of","y.",":","return",":","tuple","containing","batch","data","loss",",","batch","data","accuracy","and","number","of","samples","in","the","batch"]
code: def _batch_process(
      self,
      i_epoch: int,
      nb_epochs: int,
      batch_id: int,
      x_batch: np.ndarray,
      y_batch: np.ndarray,
  ) -> Tuple[float, float, float]:
      """
      Perform the operations of OAAT for a batch of data.
      See class documentation for more information on the exact procedure.
  
      :param i_epoch: training epoch number.
      :param nb_epochs: total training epochs.
      :param batch_id: batch_id of input data.
      :param x_batch: batch of x.
      :param y_batch: batch of y.
      :return: tuple containing batch data loss, batch data accuracy and number of samples in the batch
      """
      import torch
      from torch import nn
      import torch.nn.functional as F
  
      if self._classifier.optimizer is None:
          raise ValueError(
              "Optimizer of classifier is currently None, but is required for adversarial training."
          )
  
      if self._proxy_classifier.optimizer is None:
          raise ValueError(
              "Optimizer of proxy classifier is currently None, but is required for adversarial training."
          )
  
      self._classifier.model.train(mode=False)
  
      if self._apply_lpips_pert and (batch_id % 2 != 0):
          x_batch_pert = self._attack_lpips(
              x_batch,
              y=y_batch,
              eps=self._train_params["i_epsilon"],
              eps_step=self._train_params["i_epsilon"] / 4.0,
              max_iter=self._train_params["max_iter"],
              training_mode=False,
          )
  
      elif self._apply_lpips_pert and (batch_id % 2 == 0):
          self._attack.set_params(
              eps=self._train_params["mixup_epsilon"],
              eps_step=self._train_params["mixup_epsilon"] / 4.0,
              max_iter=self._train_params["max_iter"],
          )
          x_batch_pert = self._attack.generate(x_batch, y=y_batch)
  
      else:
          if i_epoch <= (nb_epochs // 4):
              self._attack.set_params(
                  eps=self._train_params["i_epsilon"],
                  eps_step=self._train_params["i_epsilon"] / 2.0,
                  max_iter=5,
              )
          else:
              self._attack.set_params(
                  eps=self._train_params["i_epsilon"],
                  eps_step=self._train_params["i_epsilon"] / 4.0,
                  max_iter=self._train_params["max_iter"],
              )
  
          x_batch_pert = self._attack.generate(x_batch, y=y_batch)
  
      # Apply preprocessing
      y_batch = check_and_transform_label_format(
          y_batch, nb_classes=self.classifier.nb_classes
      )
  
      (
          x_preprocessed,
          y_preprocessed,
      ) = self._classifier._apply_preprocessing(  # pylint: disable=W0212
          x_batch, y_batch, fit=True
      )
      (
          x_preprocessed_pert,
          _,
      ) = self._classifier._apply_preprocessing(  # pylint: disable=W0212
          x_batch_pert, y_batch, fit=True
      )
  
      # Check label shape
      if self._classifier._reduce_labels:  # pylint: disable=W0212
          y_preprocessed = np.argmax(y_preprocessed, axis=1)
  
      i_batch = torch.from_numpy(x_preprocessed).to(
          self._classifier.device
      )
      i_batch_pert = torch.from_numpy(x_preprocessed_pert).to(
          self._classifier.device
      )
      o_batch = torch.from_numpy(y_preprocessed).to(
          self._classifier.device
      )
  
      self._classifier.model.train(mode=True)
  
      if self._apply_wp:
          if not self._apply_lpips_pert:
              i_batch_pert_awp = torch.clamp(
                  i_batch + 2.0 * (i_batch_pert - i_batch), 0, 1
              )
          else:
              i_batch_pert_awp = i_batch_pert
          w_perturb = self._weight_perturbation(
              x_batch=i_batch,
              x_batch_pert=i_batch_pert_awp,
              y_batch=o_batch,
          )
          list_keys = list(w_perturb.keys())
          self._modify_classifier(
              self._classifier, list_keys, w_perturb, op="add"
          )
  
      # Zero the parameter gradients
      self._classifier.optimizer.zero_grad()
  
      n = x_batch.shape[0]
      # Perform prediction
      model_outputs = self._classifier.model(i_batch)
      model_outputs_pert = self._classifier.model(i_batch_pert)
  
      # Form the loss function
      loss_clean = self._classifier.loss(model_outputs, o_batch)
  
      if self._apply_lpips_pert and (batch_id % 2 == 0):
          i_batch_pert_eps = torch.min(
              torch.max(
                  i_batch_pert,
                  i_batch - self._train_params["i_epsilon"],
              ),
              i_batch + self._train_params["i_epsilon"],
          )
          model_outputs_pert_eps = self._classifier.model(
              i_batch_pert_eps
          )
          loss_kl = (1.0 / n) * nn.KLDivLoss(reduction="sum")(
              F.log_softmax(model_outputs_pert_eps, dim=1),
              torch.clamp(
                  self._train_params["alpha"]
                  * F.softmax(model_outputs, dim=1)
                  + (1.0 - self._train_params["alpha"])
                  * F.softmax(model_outputs_pert, dim=1),
                  min=EPS,
              ),
          )
  
      else:
          loss_kl = (1.0 / n) * nn.KLDivLoss(reduction="sum")(
              F.log_softmax(model_outputs_pert, dim=1),
              torch.clamp(F.softmax(model_outputs, dim=1), min=EPS),
          )
      loss = loss_clean + self._train_params["beta"] * loss_kl
  
      loss.backward()
  
      self._classifier.optimizer.step()
  
      if self._apply_wp:
          self._modify_classifier(
              self._classifier, list_keys, w_perturb, op="subtract"
          )
  
      train_loss = loss.item() * o_batch.size(0)
      train_acc = (model_outputs_pert.max(1)[1] == o_batch).sum().item()
      train_n = o_batch.size(0)
  
      for epoch_swa, tau_val, p_classifier in zip(
          self._train_params["list_swa_epoch"],
          self._train_params["list_swa_tau"],
          self._list_avg_models,
      ):
          if i_epoch == epoch_swa:
              for (
                  name,
                  params,
              ) in self._classifier.model.state_dict().items():
                  p_classifier.model.state_dict()[name] = params
          elif i_epoch > epoch_swa:
              for (
                  name,
                  params,
              ) in self._classifier.model.state_dict().items():
                  p_classifier.model.state_dict()[name] = (
                      (1 - tau_val) * params
                      + tau_val * p_classifier.model.state_dict()[name]
                  )
          else:
              pass
  
      self._classifier.model.train(mode=False)
  
      return train_loss, train_acc, train_n
  
imports: from __future__ import absolute_import, division, print_function, unicode_literals
  from collections import OrderedDict
  import logging
  import os
  import time
  from typing import Optional, Tuple, TYPE_CHECKING, List, Dict, Union
  import six
  import numpy
  from tqdm.auto import trange
  from art import config
  from art.defences.trainer.adversarial_trainer_oaat import AdversarialTrainerOAAT
  from art.estimators.classification.pytorch import PyTorchClassifier
  from art.data_generators import DataGenerator
  from art.attacks.attack import EvasionAttack
  from art.utils import check_and_transform_label_format, random_sphere
  from art.config import ART_NUMPY_DTYPE
label_desc: Use image node_id 4 for calling the AdversarialTrainerOAATPyTorch obj's underlying member method code with example usage: obj._batch_process(i_epoch, nb_epochs, batch_id, x_batch, y_batch) and returns: train_loss,  train_acc,  train_n