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: Train a model adversarially with OAAT protocol using a data generator.
  See class documentation for more information on the exact procedure.
  
  :param generator: Data generator.
  :param validation_data: Tuple consisting of validation data, (x_val, y_val)
  :param nb_epochs: Number of epochs to use for trainings.
  :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                            the target classifier.
docstr_tok: ["Train","a","model","adversarially","with","OAAT","protocol","using","a","data","generator",".","See","class","documentation","for","more","information","on","the","exact","procedure",".",":","param","generator",":","Data","generator",".",":","param","validation_data",":","Tuple","consisting","of","validation","data",",","(","x_val",",","y_val",")",":","param","nb_epochs",":","Number","of","epochs","to","use","for","trainings",".",":","param","kwargs",":","Dictionary","of","framework-specific","arguments",".","These","will","be","passed","as","such","to","the","`","fit","`","function","of","the","target","classifier","."]
code: def fit_generator(
      self,
      generator: DataGenerator,
      validation_data: Optional[Tuple[np.ndarray, np.ndarray]] = None,
      nb_epochs: int = 20,
      **kwargs,
  ):  # pylint: disable=W0221
      """
      Train a model adversarially with OAAT protocol using a data generator.
      See class documentation for more information on the exact procedure.
  
      :param generator: Data generator.
      :param validation_data: Tuple consisting of validation data, (x_val, y_val)
      :param nb_epochs: Number of epochs to use for trainings.
      :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                                the target classifier.
      """
      import torch
  
      logger.info("Performing adversarial training with OAAT protocol")
  
      size = generator.size
      batch_size = generator.batch_size
      if size is not None:
          nb_batches = int(np.ceil(size / batch_size))
      else:
          raise ValueError("Size is None.")
  
      logger.info("Adversarial Training OAAT")
  
      if (
          len(self._train_params["list_swa_epoch"])
          != len(self._train_params["list_swa_tau"])
      ) or (
          len(self._train_params["list_swa_epoch"])
          != len(self._list_avg_models)
      ):
          raise ValueError(
              "number of elements of list_swa_epoch must be same as of  "
              "list_swa_tau and list_avg_models."
          )
  
      beta_init = self._train_params["beta"]
      best_acc_adv_test = 0
      for i_epoch in trange(
          nb_epochs, desc="Adversarial Training OAAT - Epochs"
      ):
          if i_epoch >= self._train_params["awp_warmup"]:
              self._apply_wp = True
  
          if i_epoch == (int(3 * nb_epochs / 4)):
              self._classifier._optimizer = (
                  torch.optim.SGD(  # pylint: disable=W0212
                      self._classifier.model.parameters(),
                      lr=self._train_params["lr"],
                      momentum=self._train_params["momentum"],
                      weight_decay=self._train_params["weight_decay"],
                  )
              )
  
          if isinstance(self._classifier.optimizer, torch.optim.SGD):
              if self._train_params["oaat_warmup"] == 1:
                  if i_epoch < 10:
                      self.update_learning_rate(
                          self._classifier.optimizer,
                          i_epoch,
                          nb_epochs,
                          lr_schedule="linear",
                      )
                  else:
                      self.update_learning_rate(
                          self._classifier.optimizer,
                          i_epoch,
                          nb_epochs,
                          lr_schedule=self._train_params["lr_schedule"],
                      )
              else:
                  self.update_learning_rate(
                      self._classifier.optimizer,
                      i_epoch,
                      nb_epochs,
                      lr_schedule=self._train_params["lr_schedule"],
                  )
  
          if i_epoch > (nb_epochs // 4):
              self._train_params["i_epsilon"] = (
                  i_epoch * self._train_params["epsilon"] / nb_epochs
              )
  
          if (
              self._train_params["i_epsilon"]
              > self._train_params["alternate_iter_eps"]
          ):
              self._apply_lpips_pert = True
  
              load_swa_model_tau = self._train_params[
                  "load_swa_model_tau"
              ]
              file_name = f"oaat_swa_tau_{load_swa_model_tau}.model"
              if self._train_params["models_path"] is None:
                  full_path_load = os.path.join(
                      config.ART_DATA_PATH, file_name
                  )
              else:
                  full_path_load = os.path.join(
                      self._train_params["models_path"], file_name
                  )
              if os.path.isfile(full_path_load):
                  self._lpips_classifier._model._model.load_state_dict(  # pylint: disable=W0212
                      torch.load(full_path_load)
                  )
              else:
                  raise ValueError(
                      "Invalid path/file for weight average model is provided for adversarial training."
                  )
  
              self._lpips_classifier.model.train(mode=False)
              self._lpips_classifier.model.to(self._classifier.device)
  
              self._train_params["alpha"] = self._train_params[
                  "alpha"
              ] - self._train_params["mixup_alpha"] / (
                  nb_epochs - int(3 * nb_epochs / 4) + 1
              )
              self._train_params["i_lpips_weight"] = (
                  self._train_params["lpips_weight"]
                  * (i_epoch - int(3 * nb_epochs / 4))
                  / (nb_epochs - int(3 * nb_epochs / 4))
              )
              self._train_params["beta"] = beta_init + (
                  self._train_params["beta_final"] - beta_init
              ) * (i_epoch - int(3 * nb_epochs / 4)) / (
                  nb_epochs - int(3 * nb_epochs / 4)
              )
  
          start_time = time.time()
          train_loss = 0.0
          train_acc = 0.0
          train_n = 0.0
  
          for batch_id in range(nb_batches):
              # Create batch data
              x_batch, y_batch = generator.get_batch()
              x_batch = x_batch.copy()
  
              _train_loss, _train_acc, _train_n = self._batch_process(
                  i_epoch, nb_epochs, batch_id, x_batch, y_batch
              )
  
              train_loss += _train_loss
              train_acc += _train_acc
              train_n += _train_n
  
          train_time = time.time()
  
          # compute accuracy
          if validation_data is not None:
              (x_test, y_test) = validation_data
              y_test = check_and_transform_label_format(
                  y_test, nb_classes=self.classifier.nb_classes
              )
              # pylint: disable=W0212
              (
                  x_preprocessed_test,
                  y_preprocessed_test,
              ) = self._classifier._apply_preprocessing(
                  x_test,
                  y_test,
                  fit=True,
              )
              # pylint: enable=W0212
              output_clean = np.argmax(
                  self.predict(x_preprocessed_test), axis=1
              )
              nb_correct_clean = np.sum(
                  output_clean == np.argmax(y_preprocessed_test, axis=1)
              )
              self._attack.set_params(
                  eps=self._train_params["epsilon"],
                  eps_step=self._train_params["epsilon"] / 4.0,
                  max_iter=self._train_params["max_iter"],
              )
              x_test_adv = self._attack.generate(
                  x_preprocessed_test, y=y_preprocessed_test
              )
              # pylint: disable=W0212
              (
                  x_preprocessed_test_adv,
                  y_preprocessed_test,
              ) = self._classifier._apply_preprocessing(
                  x_test_adv,
                  y_test,
                  fit=True,
              )
              # pylint: enable=W0212
              output_adv = np.argmax(
                  self.predict(x_preprocessed_test_adv), axis=1
              )
              nb_correct_adv = np.sum(
                  output_adv == np.argmax(y_preprocessed_test, axis=1)
              )
  
              logger.info(
                  "epoch: %s time(s): %.1f loss: %.4f acc-adv (tr): %.4f acc-clean (val): %.4f acc-adv (val): %.4f",
                  i_epoch,
                  train_time - start_time,
                  train_loss / train_n,
                  train_acc / train_n,
                  nb_correct_clean / x_test.shape[0],
                  nb_correct_adv / x_test.shape[0],
              )
              # save last checkpoint
              if i_epoch + 1 == nb_epochs:
                  self._classifier.save(
                      filename=f"oaat_epoch_{i_epoch}",
                      path=self._train_params["models_path"],
                  )
  
              # save best checkpoint
              if nb_correct_adv / x_test.shape[0] > best_acc_adv_test:
                  self._classifier.save(
                      filename="oaat_epoch_best",
                      path=self._train_params["models_path"],
                  )
                  best_acc_adv_test = nb_correct_adv / x_test.shape[0]
  
          else:
              logger.info(
                  "epoch: %s time(s): %.1f loss: %.4f acc-adv: %.4f",
                  i_epoch,
                  train_time - start_time,
                  train_loss / train_n,
                  train_acc / train_n,
              )
  
          if i_epoch > self._train_params["swa_save_epoch"] - 1:
              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 epoch_swa <= i_epoch:
                      file_name = f"oaat_swa_tau_{tau_val}.model"
                      if self._train_params["models_path"] is None:
                          full_path = os.path.join(
                              config.ART_DATA_PATH, file_name
                          )
                      else:
                          full_path = os.path.join(
                              self._train_params["models_path"],
                              file_name,
                          )
  
                      folder = os.path.split(full_path)[0]
                      if not os.path.exists(folder):
                          os.makedirs(folder)
                      # pylint: disable=W0212
                      # disable pylint because access to _modules required
                      torch.save(
                          p_classifier._model._model.state_dict(),
                          full_path,
                      )
  
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 3 for calling the AdversarialTrainerOAATPyTorch obj's underlying member method code with example usage: obj.fit_generator(generator, validation_data, nb_epochs) without return types