class_docstr: Implementation of SmoothMix training, as introduced in Jeong et al. (2021)
  
  | Paper link: https://arxiv.org/abs/2111.09277
class_docstr_tok: ["Implementation","of","SmoothMix","training",",","as","introduced","in","Jeong","et","al",".","(","2021",")","|","Paper","link",":","https",":","\/\/arxiv.org\/abs\/2111.09277"]
docstr: Fit the classifier on the training set `(x, y)`.
  
  :param x: Training data.
  :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
            shape (nb_samples,).
  :param batch_size: Size of batches.
  :param nb_epochs: Number of epochs to use for training.
  :param training_mode: `True` for model set to training mode and `'False` for model set to evaluation mode.
  :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                    the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                    the last batch will be smaller. (default: ``False``)
  :param scheduler: Learning rate scheduler to run at the start of every epoch.
  :param verbose: Display the training progress bar.
  :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
         and providing it takes no effect.
docstr_tok: ["Fit","the","classifier","on","the","training","set","`","(","x",",","y",")","`",".",":","param","x",":","Training","data",".",":","param","y",":","Target","values","(","class","labels",")","one-hot-encoded","of","shape","(","nb_samples",",","nb_classes",")","or","index","labels","of","shape","(","nb_samples",",",")",".",":","param","batch_size",":","Size","of","batches",".",":","param","nb_epochs",":","Number","of","epochs","to","use","for","training",".",":","param","training_mode",":","`","True","`","for","model","set","to","training","mode","and","`","'False","`","for","model","set","to","evaluation","mode",".",":","param","drop_last",":","Set","to","``","True","``","to","drop","the","last","incomplete","batch",",","if","the","dataset","size","is","not","divisible","by","the","batch","size",".","If","``","False","``","and","the","size","of","dataset","is","not","divisible","by","the","batch","size",",","then","the","last","batch","will","be","smaller",".","(","default",":","``","False","``",")",":","param","scheduler",":","Learning","rate","scheduler","to","run","at","the","start","of","every","epoch",".",":","param","verbose",":","Display","the","training","progress","bar",".",":","param","kwargs",":","Dictionary","of","framework-specific","arguments",".","This","parameter","is","not","currently","supported","for","PyTorch","and","providing","it","takes","no","effect","."]
code: def fit(  # pylint: disable=W0221
      self,
      x: np.ndarray,
      y: np.ndarray,
      batch_size: int = 128,
      nb_epochs: int = 10,
      training_mode: bool = True,
      drop_last: bool = False,
      scheduler: Optional[
          "torch.optim.lr_scheduler._LRScheduler"
      ] = None,
      verbose: bool = False,
      **kwargs,
  ) -> None:
      """
      Fit the classifier on the training set `(x, y)`.
  
      :param x: Training data.
      :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                shape (nb_samples,).
      :param batch_size: Size of batches.
      :param nb_epochs: Number of epochs to use for training.
      :param training_mode: `True` for model set to training mode and `'False` for model set to evaluation mode.
      :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                        the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                        the last batch will be smaller. (default: ``False``)
      :param scheduler: Learning rate scheduler to run at the start of every epoch.
      :param verbose: Display the training progress bar.
      :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
             and providing it takes no effect.
      """
      import torch
      import torch.nn.functional as F
      from torch.utils.data import TensorDataset, DataLoader
  
      # Set model mode
      self._model.train(mode=training_mode)
  
      if self._optimizer is None:  # pragma: no cover
          raise ValueError(
              "An optimizer is needed to train the model, but none for provided"
          )
  
      y = check_and_transform_label_format(
          y, nb_classes=self.nb_classes
      )
  
      # Apply preprocessing
      x_preprocessed, y_preprocessed = self._apply_preprocessing(
          x, y, fit=True
      )
  
      # Check label shape
      y_preprocessed = self.reduce_labels(y_preprocessed)
  
      # Create dataloader
      x_tensor = torch.from_numpy(x_preprocessed)
      y_tensor = torch.from_numpy(y_preprocessed)
      dataset = TensorDataset(x_tensor, y_tensor)
      dataloader = DataLoader(
          dataset=dataset,
          batch_size=batch_size,
          shuffle=True,
          drop_last=drop_last,
      )
  
      # Start training
      for epoch in trange(nb_epochs, disable=not verbose):
          warmup_v = min(1.0, (epoch + 1) / self.warmup)
  
          for x_batch, y_batch in dataloader:
              mini_batch_size = len(x_batch) // self.num_noise_vec
  
              for mini_batch in range(self.num_noise_vec):
                  # Create mini batch
                  inputs = x_batch[
                      mini_batch
                      * mini_batch_size : (mini_batch + 1)
                      * mini_batch_size
                  ]
                  labels = y_batch[
                      mini_batch
                      * mini_batch_size : (mini_batch + 1)
                      * mini_batch_size
                  ]
  
                  # Move inputs to GPU
                  inputs = inputs.to(self.device)
                  labels = labels.to(self.device)
  
                  noises = [
                      torch.randn_like(inputs) * self.scale
                      for _ in range(self.num_noise_vec)
                  ]
  
                  # Attack and find adversarial examples
                  self._model.eval()
                  inputs, inputs_adv = self._smooth_mix_pgd_attack(
                      inputs, labels, noises, warmup_v
                  )
                  self._model.train(mode=training_mode)
  
                  in_clean_c = torch.cat(
                      [inputs + noise for noise in noises], dim=0
                  )
                  logits_c = self._model(in_clean_c)[-1]
                  labels_c = labels.repeat(self.num_noise_vec)
  
                  logits_c_chunk = torch.chunk(
                      logits_c, self.num_noise_vec, dim=0
                  )
                  clean_sm = F.softmax(
                      torch.stack(logits_c_chunk), dim=-1
                  )
                  clean_avg_sm = torch.mean(clean_sm, dim=0)
                  loss_xent = F.cross_entropy(
                      logits_c, labels_c, reduction="none"
                  )
  
                  # mix adversarial examples
                  in_mix, labels_mix = self._mix_data(
                      inputs, inputs_adv, clean_avg_sm
                  )
                  in_mix_c = torch.cat(
                      [in_mix + noise for noise in noises], dim=0
                  )
                  labels_mix_c = labels_mix.repeat(
                      self.num_noise_vec, 1
                  )
                  logits_mix_c = F.log_softmax(
                      self._model(in_mix_c)[-1], dim=1
                  )
  
                  preds = torch.argmax(clean_avg_sm, dim=-1)
                  ind_correct = (preds == labels).float()
                  ind_correct = ind_correct.repeat(self.num_noise_vec)
  
                  loss_mixup = F.kl_div(
                      logits_mix_c, labels_mix_c, reduction="none"
                  ).sum(1)
                  loss = (
                      loss_xent.mean()
                      + self.eta
                      * warmup_v
                      * (ind_correct * loss_mixup).mean()
                  )
  
                  # compute gradient and do SGD step
                  self._optimizer.zero_grad()
                  loss.backward()
                  self._optimizer.step()
  
          if scheduler is not None:
              scheduler.step()
  
file_path: adversarial-robustness-toolbox/art/estimators/certification/randomized_smoothing/smooth_mix/pytorch.py
imports: from __future__ import absolute_import, division, print_function, unicode_literals
  import logging
  from typing import List, Optional, Tuple, Union, TYPE_CHECKING
  from tqdm.auto import trange
  import numpy
  from art.estimators.certification.randomized_smoothing.pytorch import PyTorchRandomizedSmoothing
  from art.utils import check_and_transform_label_format
label_desc: Use image node_id 2 for calling the PyTorchSmoothMix obj's underlying member method code with example usage: obj.fit(x, y, batch_size, nb_epochs, training_mode, drop_last, scheduler, verbose) without return types