text: from __future__ import print_function
  import os
  import argparse
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  import torchvision
  import torch.optim as optim
  from torchvision import datasets, transforms
  
  from models.wideresnet import *
  from models.resnet import *
  
  parser = argparse.ArgumentParser(description='PyTorch CIFAR TRADES Adversarial Training')
  parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                      help='input batch size for training (default: 128)')
  parser.add_argument('--test-batch-size', type=int, default=128, metavar='N',
                      help='input batch size for testing (default: 128)')
  parser.add_argument('--epochs', type=int, default=76, metavar='N',
                      help='number of epochs to train')
  parser.add_argument('--weight-decay', '--wd', default=2e-4,
                      type=float, metavar='W')
  parser.add_argument('--lr', type=float, default=0.1, metavar='LR',
                      help='learning rate')
  parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                      help='SGD momentum')
  parser.add_argument('--no-cuda', action='store_true', default=False,
                      help='disables CUDA training')
  parser.add_argument('--epsilon', default=8.0/255.0,
                      help='perturbation')
  parser.add_argument('--num-steps', type=int, default=7,
                      help='perturb number of steps')
  parser.add_argument('--step-size', type=float, default=0.3,
                      help='perturb step size')
  parser.add_argument('--seed', type=int, default=1, metavar='S',
                      help='random seed (default: 1)')
  parser.add_argument('--log-interval', type=int, default=100, metavar='N',
                      help='how many batches to wait before logging training status')
  parser.add_argument('--model-dir', default='./model-cifar-wideResNet',
                      help='directory of model for saving checkpoint')
  parser.add_argument('--save-freq', '-s', default=5, type=int, metavar='N',
                      help='save frequency')
  parser.add_argument('--lbd', type=float, default=0.01,
                      help='lambda for the entropy term')
  parser.add_argument('--dataset', type=str, default='cifar10', help='dataset')
  
  args = parser.parse_args()
  print(args)
  
  # settings
  model_dir = args.model_dir
  if not os.path.exists(model_dir):
      os.makedirs(model_dir)
  use_cuda = not args.no_cuda and torch.cuda.is_available()
  torch.manual_seed(args.seed)
  device = torch.device("cuda" if use_cuda else "cpu")
  kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {}
  
  # setup data loader
  transform_train = transforms.Compose([
      transforms.RandomCrop(32, padding=4),
      transforms.RandomHorizontalFlip(),
      transforms.ToTensor(),
  ])
  transform_test = transforms.Compose([
      transforms.ToTensor(),
  ])
  if args.dataset == 'cifar10':
      trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform_train)
      testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test)
  elif args.dataset == 'cifar100':
      trainset = torchvision.datasets.CIFAR100(root='../data', train=True, download=True, transform=transform_train)
      testset = torchvision.datasets.CIFAR100(root='../data', train=False, download=True, transform=transform_test)
  elif args.dataset == 'svhn':
      args.epsilon = 4.0 / 255.0
      trainset = torchvision.datasets.SVHN(root='../data', split='train', download=True, transform=transform_test)
      testset = torchvision.datasets.SVHN(root='../data', split='test', download=True, transform=transform_test)
  else:
      raise NotImplementedError
  
  train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, **kwargs)
  test_loader = torch.utils.data.DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, **kwargs)
  
  def adt_loss(model,
               x_natural,
               y,
               optimizer,
               learning_rate=1.0,
               epsilon=8.0/255.0,
               perturb_steps=10,
               num_samples=10,
               lbd=0.0):
  
      model.eval()
      batch_size = len(x_natural)
      # generate adversarial example
      mean = Variable(torch.zeros(x_natural.size()).cuda(), requires_grad=True)
      var = Variable(torch.zeros(x_natural.size()).cuda(), requires_grad=True)
      optimizer_adv = optim.Adam([mean, var], lr=learning_rate, betas=(0.0, 0.0))
  
      for _ in range(perturb_steps):
          for s in range(num_samples):
              adv_std = F.softplus(var)
              rand_noise = torch.randn_like(x_natural)
              adv = torch.tanh(mean + rand_noise * adv_std)
              # omit the constants in -logp
              negative_logp = (rand_noise ** 2) / 2. + (adv_std + 1e-8).log() + (1 - adv ** 2 + 1e-8).log()
              entropy = negative_logp.mean() # entropy
              x_adv = torch.clamp(x_natural + epsilon * adv, 0.0, 1.0)
  
              # minimize the negative loss
              with torch.enable_grad():
                  loss = -F.cross_entropy(model(x_adv), y) - lbd * entropy
              loss.backward(retain_graph=True if s != num_samples - 1 else False)
  
          optimizer_adv.step()
      
      x_adv = torch.clamp(x_natural + epsilon * torch.tanh(mean + F.softplus(var) * torch.randn_like(x_natural)), 0.0, 1.0)
      model.train()
      x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
      # zero gradient
      optimizer.zero_grad()
      # calculate robust loss
      logits = model(x_adv)
      loss = F.cross_entropy(logits, y)
      return loss
  
  
  def train(args, model, device, train_loader, optimizer, epoch):
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          data, target = data.to(device), target.to(device)
  
          optimizer.zero_grad()
  
          loss = adt_loss(model=model,
                          x_natural=data,
                          y=target,
                          optimizer=optimizer,
                          learning_rate=args.step_size,
                          epsilon=args.epsilon,
                          perturb_steps=args.num_steps,
                          num_samples=5,
                          lbd=args.lbd)
  
          loss.backward()
          optimizer.step()
  
          # print progress
          if batch_idx % args.log_interval == 0:
              print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                  epoch, batch_idx * len(data), len(train_loader.dataset),
                         100. * batch_idx / len(train_loader), loss.item()))
  
  
  def eval_train(model, device, train_loader):
      model.eval()
      train_loss = 0
      correct = 0
      with torch.no_grad():
          for data, target in train_loader:
              data, target = data.to(device), target.to(device)
              output = model(data)
              train_loss += F.cross_entropy(output, target, size_average=False).item()
              pred = output.max(1, keepdim=True)[1]
              correct += pred.eq(target.view_as(pred)).sum().item()
      train_loss /= len(train_loader.dataset)
      print('Training: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
          train_loss, correct, len(train_loader.dataset),
          100. * correct / len(train_loader.dataset)))
      training_accuracy = correct / len(train_loader.dataset)
      return train_loss, training_accuracy
  
  
  def eval_test(model, device, test_loader):
      model.eval()
      test_loss = 0
      correct = 0
      with torch.no_grad():
          for data, target in test_loader:
              data, target = data.to(device), target.to(device)
              output = model(data)
              test_loss += F.cross_entropy(output, target, size_average=False).item()
              pred = output.max(1, keepdim=True)[1]
              correct += pred.eq(target.view_as(pred)).sum().item()
      test_loss /= len(test_loader.dataset)
      print('Test: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
          test_loss, correct, len(test_loader.dataset),
          100. * correct / len(test_loader.dataset)))
      test_accuracy = correct / len(test_loader.dataset)
      return test_loss, test_accuracy
  
  
  def adjust_learning_rate(optimizer, epoch):
      """decrease the learning rate"""
      lr = args.lr
      if epoch >= 75:
          lr = args.lr * 0.1
      if epoch >= 90:
          lr = args.lr * 0.01
      if epoch >= 100:
          lr = args.lr * 0.001
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr
  
  
  def main():
      # init model, ResNet18() can be also used here for training
      model = WideResNet(depth=28, num_classes=100 if args.dataset == 'cifar100' else 10).to(device)
      optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
  
      for epoch in range(1, args.epochs + 1):
          # adjust learning rate for SGD
          adjust_learning_rate(optimizer, epoch)
  
          # adversarial training
          train(args, model, device, train_loader, optimizer, epoch)
  
          # evaluation on natural examples
          print('================================================================')
          eval_train(model, device, train_loader)
          eval_test(model, device, test_loader)
          print('================================================================')
  
          # save checkpoint
          if epoch % args.save_freq == 0 or epoch > 70:
              torch.save(model.state_dict(),
                         os.path.join(model_dir, 'model-wideres-epoch{}.pt'.format(epoch)))
              torch.save(optimizer.state_dict(),
                         os.path.join(model_dir, 'opt-wideres-checkpoint_epoch{}.tar'.format(epoch)))
  
  
  if __name__ == '__main__':
      main()
  