class_docstr: Implement the CutMix data augmentation defence approach in PyTorch.
  
  | Paper link: https://arxiv.org/abs/1905.04899
  
  | Please keep in mind the limitations of defences. For more information on the limitations of this defence,
      see https://arxiv.org/abs/1803.09868 . For details on how to evaluate classifier security in general, see
      https://arxiv.org/abs/1902.06705
class_docstr_tok: ["Implement","the","CutMix","data","augmentation","defence","approach","in","PyTorch",".","|","Paper","link",":","https",":","\/\/arxiv.org\/abs\/1905.04899","|","Please","keep","in","mind","the","limitations","of","defences",".","For","more","information","on","the","limitations","of","this","defence",",","see","https",":","\/\/arxiv.org\/abs\/1803.09868",".","For","details","on","how","to","evaluate","classifier","security","in","general",",","see","https",":","\/\/arxiv.org\/abs\/1902.06705"]
docstr: Apply CutMix data augmentation to sample `x`.
  
  :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
  :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
            or class indices of shape `(nb_samples,)`.
  :return: Data augmented sample. The returned labels will be probability vectors of shape
           `(nb_samples, nb_classes)`.
docstr_tok: ["Apply","CutMix","data","augmentation","to","sample","`","x","`",".",":","param","x",":","Sample","to","augment","with","shape","of","`","NCHW","`",",","`","NHWC","`",",","`","NCFHW","`","or","`","NFHWC","`",".",":","param","y",":","Labels","of","`","x","`","either","one-hot","or","multi-hot","encoded","of","shape","`","(","nb_samples",",","nb_classes",")","`","or","class","indices","of","shape","`","(","nb_samples",",",")","`",".",":","return",":","Data","augmented","sample",".","The","returned","labels","will","be","probability","vectors","of","shape","`","(","nb_samples",",","nb_classes",")","`","."]
code: def forward(
      self, x: "torch.Tensor", y: Optional["torch.Tensor"] = None
  ) -> Tuple["torch.Tensor", Optional["torch.Tensor"]]:
      """
      Apply CutMix data augmentation to sample `x`.
  
      :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
      :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
                or class indices of shape `(nb_samples,)`.
      :return: Data augmented sample. The returned labels will be probability vectors of shape
               `(nb_samples, nb_classes)`.
      """
      import torch
  
      if y is None:
          raise ValueError("Labels `y` cannot be None.")
  
      # convert labels to one-hot encoding
      if len(y.shape) == 2:
          y_one_hot = y
      elif len(y.shape) == 1:
          y_one_hot = torch.nn.functional.one_hot(y, self.num_classes)
      else:
          raise ValueError(
              "Shape of labels not recognised. "
              "Please provide labels in shape (nb_samples,) or (nb_samples, nb_classes)"
          )
  
      x_ndim = len(x.shape)
  
      # NCHW/NHWC/NFHWC --> NCFHW
      if x_ndim == 4:
          if self.channels_first:
              # NCHW --> NCFHW
              x_ncfhw = torch.unsqueeze(x, dim=2)
          else:
              # NHWC --> NCHW --> NCFHW
              x_nchw = torch.permute(x, (0, 3, 1, 2))
              x_ncfhw = torch.unsqueeze(x_nchw, dim=2)
      elif x_ndim == 5:
          if self.channels_first:
              # NCFHW
              x_ncfhw = x
          else:
              # NFHWC --> NCFHW
              x_ncfhw = torch.permute(x, (0, 4, 1, 2, 3))
      else:
          raise ValueError(
              "Unrecognized input dimension. CutMix can only be applied to image and video data."
          )
  
      n, _, _, height, width = x_ncfhw.shape
      x_aug = x_ncfhw.clone()
      y_aug = y_one_hot.clone().float()
  
      # sample the combination ratio from the Beta distribution
      lmb = np.random.beta(self.alpha, self.alpha)
      cut_ratio = np.sqrt(1.0 - lmb)
      cut_height = int(height * cut_ratio)
      cut_width = int(width * cut_ratio)
  
      # randomly choose indices for samples to mix
      indices = torch.randperm(n, device=x.device)
  
      # generate a random bounding box per image
      for idx1, idx2 in enumerate(
          tqdm(indices, desc="CutMix", disable=not self.verbose)
      ):
          prob = np.random.rand()
          if prob < self.probability:
              # uniform sampling
              center_x = torch.randint(0, height, (1,))
              center_y = torch.randint(0, width, (1,))
              bby1 = torch.clamp(center_y - cut_height // 2, 0, height)
              bbx1 = torch.clamp(center_x - cut_width // 2, 0, width)
              bby2 = torch.clamp(center_y + cut_height // 2, 0, height)
              bbx2 = torch.clamp(center_x + cut_width // 2, 0, width)
  
              # insert image bounding box
              x_aug[idx1, :, :, bbx1:bbx2, bby1:bby2] = x_ncfhw[idx2, :, :, bbx1:bbx2, bby1:bby2]  # type: ignore
              # mix labels
              y_aug[idx1] = lmb * y_aug[idx1] + (1.0 - lmb) * y_one_hot[idx2]  # type: ignore
  
      x_ncfhw = x_aug
  
      # NCHW/NHWC/NFHWC <-- NCFHW
      if x_ndim == 4:
          if self.channels_first:
              # NCHW <-- NCFHW
              x_aug = torch.squeeze(x_ncfhw, dim=2)
          else:
              # NHWC <-- NCHW <-- NCFHW
              x_nchw = torch.squeeze(x_ncfhw, dim=2)
              x_aug = torch.permute(x_nchw, (0, 2, 3, 1))
      elif x_ndim == 5:
          if self.channels_first:
              # NCFHW
              x_aug = x_ncfhw
          else:
              # NFHWC <-- NCFHW
              x_aug = torch.permute(x_ncfhw, (0, 2, 3, 4, 1))
  
      return x_aug, y_aug
  
file_path: adversarial-robustness-toolbox/art/defences/preprocessor/cutmix/cutmix_pytorch.py
imports: from __future__ import absolute_import, division, print_function, unicode_literals
  import logging
  from typing import Optional, Tuple, TYPE_CHECKING
  import numpy
  from tqdm.auto import tqdm
  from art.defences.preprocessor.preprocessor import PreprocessorPyTorch
label_desc: Use image node_id 2 for calling the CutMixPyTorch obj's underlying member method code with example usage: obj.forward(x, y) and returns: x_aug,  y_aug