text: 
  from model import model_list, normalization, input_shape
  from model import mnist_model_toy, cifar_model_toy, mnist_model_tiny, cifar_model_tiny
  from utils import data_loaders
  
  
  import argparse
  import os
  import numpy as np
  
  import torch
  import torch.optim as optim
  from torch import nn
  
  name_mapping = {
      ('mnist', 'toy', 'random', None): 'mnist.toy.random',
      ('mnist', 'toy', 'clean', None): 'mnist.toy.clean',
      ('mnist', 'toy', 'pgdadv', '0.1'): 'mnist.toy.pgdadv.0.1',
      ('mnist', 'toy', 'pgdadv', '0.3'): 'mnist.toy.pgdadv.0.3',
      ('cifar', 'toy', 'random', None): 'cifar.toy.random',
      ('cifar', 'toy', 'clean', None): 'cifar.toy.clean',
      ('cifar', 'toy', 'pgdadv', '0.0349'): 'cifar.toy.pgdadv.2',
      ('cifar', 'toy', 'pgdadv', '0.1394'): 'cifar.toy.pgdadv.8',
      ('mnist', 'tiny', 'random', None): 'mnist.tiny.random',
      ('mnist', 'tiny', 'clean', None): 'mnist.tiny.clean',
      ('cifar', 'tiny', 'random', None): 'cifar.tiny.random',
      ('cifar', 'tiny', 'clean', None): 'cifar.tiny.clean'
  }
  
  
  def load_model(dataset, model_size):
      if model_size == 'toy':
          if dataset == 'mnist':
              return mnist_model_toy()
          elif dataset == 'cifar':
              return cifar_model_toy()
      elif model_size == 'tiny':
          if dataset == 'mnist':
              return mnist_model_tiny()
          elif dataset == 'cifar':
              return cifar_model_tiny()
      raise NotImplementedError()
  
  
  """
      adapted from https://github.com/wanglouis49/pytorch-adversarial_box/blob/master/adversarialbox/attacks.py
  """
  def to_var(x, requires_grad=False):
      """
      Varialbe type that automatically choose cpu or cuda
      """
      if torch.cuda.is_available():
          x = x.cuda()
      return torch.tensor(x, requires_grad=requires_grad)
  
  
  class LinfPGDAttack(object):
      def __init__(self, model=None, epsilon=0.3, k=40, a=0.01, random_start=True):
          """
              Attack parameter initialization. The attack performs k steps of
              size a, while always staying within epsilon from the initial
              point.
              https://github.com/MadryLab/mnist_challenge/blob/master/pgd_attack.py
  
              NOTE: the epsilon has not been rescaled
          """
          self.model = model
          self.epsilon = epsilon
          self.k = k
          self.a = a
          self.rand = random_start
          self.loss_fn = nn.CrossEntropyLoss()
  
      def perturb(self, X_nat, y, in_min=0., in_max=1.):
          """
          Given examples (X_nat, y), returns adversarial
          examples within epsilon of X_nat in l_infinity norm.
          """
          if self.rand:
              X = X_nat + np.random.uniform(-self.epsilon, self.epsilon,
                  X_nat.shape).astype('float32')
          else:
              X = np.copy(X_nat)
  
          for i in range(self.k):
              X_var = to_var(torch.from_numpy(X), requires_grad=True)
              y_var = to_var(torch.LongTensor(y))
  
              scores = self.model(X_var)
              loss = self.loss_fn(scores, y_var)
              loss.backward()
              grad = X_var.grad.data.cpu().numpy()
  
              X += self.a * np.sign(grad)
  
              X = np.clip(X, X_nat - self.epsilon, X_nat + self.epsilon)
              X = np.clip(X, in_min, in_max) # ensure valid pixel range
  
          return X
  
  
  parser = argparse.ArgumentParser(description='Calculate model adversarial lower bound by PGD')
  parser.add_argument('--cifar', default=False, action='store_true')
  parser.add_argument('--mnist', default=False, action='store_true')
  parser.add_argument('--toy', default=False, action='store_true')
  parser.add_argument('--tiny', default=False, action='store_true')
  parser.add_argument('--random', default=False, action='store_true')
  parser.add_argument('--clean', default=False, action='store_true')
  parser.add_argument('--pgdadv', default=False, action='store_true')
  parser.add_argument('--eps')
  parser.add_argument('--cuda_ids', type=int, default=0)
  parser.add_argument('--batch_size', type=int, default=512)
  parser.add_argument('--epochs', type=int, default=20)
  parser.add_argument('--lr', type=float, default=0.1)
  parser.add_argument('--momentum', type=float, default=0.1)
  parser.add_argument('--weight_decay', type=float, default=5e-4)
  parser.add_argument('--pgd_step', type=int, default=20)
  if __name__ == '__main__':
      args = parser.parse_args()
  
      os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_ids)
  
      tasklist = list()
      for dataset in ['mnist', 'cifar']:
          cont = (dataset == 'mnist' and args.mnist) or (dataset == 'cifar' and args.cifar)
          if cont:
              for model_size in ['toy', 'tiny']:
                  cont = (model_size == 'toy' and args.toy) or (model_size == 'tiny' and args.tiny)
                  if cont:
                      for mode in ['random', 'clean', 'pgdadv']:
                          cont = (mode == 'random' and args.random) or (mode == 'clean' and args.clean) or (mode == 'pgdadv' and args.pgdadv and args.eps is not None)
                          if cont:
                              tasklist.append((dataset, model_size, mode, args.eps if mode == 'pgdadv' else None))
  
      for task in tasklist:
          dataset, model_size, mode, eps = task
          model = load_model(dataset, model_size)
  
          if dataset == 'mnist':
              mean = normalization['mnist'][0]
              std = [normalization['mnist'][1] for _ in mean]
          elif dataset == 'cifar':
              mean = normalization['cifar'][0]
              std = [normalization['cifar'][1] for _ in mean]
          in_min, in_max = (0.0 - max(mean)) / min(std), (1.0 - min(mean)) / max(std)
          train_loader, test_loader = data_loaders(dataset, args.batch_size,
                                                   shuffle_test=False, norm_mean=mean, norm_std=std)
  
          if mode == 'random':
              # direct save :)
              model = model.cuda()
              torch.save(model.state_dict(), model_list[name_mapping[task]][2])
          elif mode == 'clean' or mode == 'pgdadv':
              # direct train :)
              # pgd train :)
  
              model = model.cuda()
              now_lr = args.lr
  
              step = 10
  
              if mode == 'pgdadv':
                  attacker = LinfPGDAttack()
                  eps = float(args.eps)
  
              for ep in range(args.epochs):
                  print(f'epoch {ep}')
                  if ep % 10 == 0 and ep > 0:
                      # slow it down a bit
                      now_lr /= 3.
                  opt = optim.SGD(model.parameters(), lr=now_lr, momentum=args.momentum, weight_decay=args.weight_decay)
                  tot, tot_correct, now_step = 0, 0, 0
                  for X, y in train_loader:
                      if mode == 'clean':
                          pass
                      else:
                          # apply attack
                          X, y = X.numpy(), y.numpy()
                          attacker = LinfPGDAttack(model, eps, k=args.pgd_step, a=eps / args.pgd_step * 2.)
                          X = attacker.perturb(X, y, in_min, in_max)
                          X, y = torch.tensor(X), torch.tensor(y)
                      X, y = X.cuda(), y.cuda()
  
                      out = model(X)
                      err = nn.CrossEntropyLoss()(out, y)
  
                      loss = err.item()
                      acc = ((torch.argmax(out, dim=1)) == y).sum().item()
  
                      opt.zero_grad()
                      err.backward()
                      opt.step()
  
                      tot += X.shape[0]
                      tot_correct += acc
                      now_step += 1
                      if now_step >= step:
                          now_step = 0
                          print(f'  epoch {ep}, samples {tot}, loss {loss:.3}, now acc {acc / X.shape[0]:.3}, tot acc {tot_correct / tot:.3}')
  
                  print('test set')
                  tot, tot_correct, now_step = 0, 0, 0
                  best = 0.0
                  for X, y in test_loader:
                      if mode == 'clean':
                          pass
                      else:
                          # apply attack
                          X, y = X.numpy(), y.numpy()
                          attacker = LinfPGDAttack(model, eps, k=args.pgd_step, a=eps / args.pgd_step * 2.)
                          X = attacker.perturb(X, y, in_min, in_max)
                          X, y = torch.tensor(X), torch.tensor(y)
                      X, y = X.cuda(), y.cuda()
  
                      with torch.no_grad():
                          out = model(X)
                          acc = ((torch.argmax(out, dim=1)) == y).sum().item()
  
                      tot += X.shape[0]
                      tot_correct += acc
                      now_step += 1
                      if now_step >= step:
                          now_step = 0
                          print(f'  epoch {ep}, samples {tot}, tot acc {tot_correct / tot:.3}')
                  if tot_correct / tot > best:
                      print(f'save with {tot_correct / tot:.3}')
                      torch.save(model.state_dict(), model_list[name_mapping[task]][2])
  
          print(f'{task} done')
  