text: '''Train CIFAR10 with PyTorch.'''
  from __future__ import print_function
  
  import torch
  import torch.nn as nn
  import torch.optim as optim
  import torch.nn.functional as F
  import torch.backends.cudnn as cudnn
  from torch.optim.lr_scheduler import CosineAnnealingLR
  
  import torchvision
  import torchvision.transforms as transforms
  
  import os
  import argparse
  import csv
  
  import models
  
  from utils import progress_bar, mixup_data, mixup_criterion
  
  import numpy
  import random
  
  model_names = sorted(name for name in models.__dict__
      if name.islower() and not name.startswith("__")
      and callable(models.__dict__[name]))
  
  parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
  parser.add_argument('-a', '--arch', metavar='ARCH', default='fixup_resnet110',
                      choices=model_names, help='model architecture: ' +
                          ' | '.join(model_names) + ' (default: fixup_resnet110)')
  parser.add_argument("--resume", default="", type=str,
                      help="path to latest checkpoint (default: none)")
  parser.add_argument('--dataroot', help='path to dataset',
                      default='/scratch/ssd/data')
  parser.add_argument('--logdir', help="path to store checkpoints")
                      #default='/scratch/ssd/logs/bn-robust/cifar10/zhang-fixup')
  parser.add_argument('--sess', default='mixup_default', type=str, help='session id')
  parser.add_argument('--seed', default=0, type=int, help='rng seed')
  parser.add_argument('--alpha', default=0., type=float, help='interpolation strength (uniform=1., ERM=0.)')
  parser.add_argument('--sgdr', action='store_true', help='use SGD with cosine annealing learning rate and restarts')
  parser.add_argument('--decay', default=5e-4, type=float, help='weight decay (default=1e-4)')
  parser.add_argument('--batchsize', default=128, type=int, help='batch size per GPU (default=128)')
  parser.add_argument('--n_epoch', default=200, type=int, help='total number of epochs')
  parser.add_argument('--base_lr', default=0.1, type=float, help='base learning rate (default=0.1)')
  
  args = parser.parse_args()
  
  torch.manual_seed(args.seed)
  torch.cuda.manual_seed(args.seed)
  numpy.random.seed(args.seed)
  random.seed(args.seed)
  torch.backends.cudnn.deterministic = True
  
  use_cuda = torch.cuda.is_available()
  best_acc = 0  # best test accuracy
  start_epoch = 0  # start from epoch 0 or last checkpoint epoch
  batch_size = args.batchsize
  base_learning_rate = args.base_lr * args.batchsize / 128.
  if use_cuda:
      # data parallel
      n_gpu = torch.cuda.device_count()
      batch_size *= n_gpu
      base_learning_rate *= n_gpu
  
  # Data
  print('==> Preparing data..')
  transform_train = transforms.Compose([
      transforms.RandomCrop(32, padding=4),
      transforms.RandomHorizontalFlip(),
      transforms.ToTensor(),
      transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
  ])
  
  transform_test = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
  ])
  
  trainset = torchvision.datasets.CIFAR10(root=args.dataroot, train=True,
                                          download=True, transform=transform_train)
  trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                            shuffle=True, num_workers=2)
  
  testset = torchvision.datasets.CIFAR10(root=args.dataroot, train=False,
                                         download=True, transform=transform_test)
  testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                           shuffle=False, num_workers=2)
  
  classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
  
  # Model
  if args.resume:
      if os.path.isfile(args.resume):
          # Load checkpoint.
          print('==> Resuming from checkpoint..')
          #checkpoint_file = os.path.join(args.resume, 'checkpoint/ckpt.t7.') + \
          #                  args.sess + '_' + str(args.seed)
          checkpoint = torch.load(args.resume)
          net = checkpoint['net']
          best_acc = checkpoint['acc']
          start_epoch = checkpoint['epoch'] + 1
          torch.set_rng_state(checkpoint['rng_state'])
  else:
      print("=> creating model '{}'".format(args.arch))
      net = models.__dict__[args.arch]()
  
  result_folder = os.path.join(args.logdir, 'results/')
  if not os.path.exists(result_folder):
      os.makedirs(result_folder)
  
  logname = result_folder + args.arch + '_' + args.sess + '_' + str(args.seed) + '.csv'
  
  if use_cuda:
      net.cuda()
      net = torch.nn.DataParallel(net)
      print('Using', torch.cuda.device_count(), 'GPUs.')
      cudnn.benchmark = True
      print('Using CUDA..')
  
  cel = nn.CrossEntropyLoss()
  criterion = lambda pred, target, lam: (-F.log_softmax(pred, dim=1) * torch.zeros(pred.size()).cuda().scatter_(1, target.data.view(-1, 1), lam.view(-1, 1))).sum(dim=1).mean()
  parameters_bias = [p[1] for p in net.named_parameters() if 'bias' in p[0]]
  parameters_scale = [p[1] for p in net.named_parameters() if 'scale' in p[0]]
  parameters_others = [p[1] for p in net.named_parameters() if not ('bias' in p[0] or 'scale' in p[0])]
  optimizer = optim.SGD(
          [{'params': parameters_bias, 'lr': args.base_lr/10.},
          {'params': parameters_scale, 'lr': args.base_lr/10.},
          {'params': parameters_others}],
          lr=base_learning_rate,
          momentum=0.9,
          weight_decay=args.decay)
  
  # Training
  def train(epoch):
      print('\nEpoch: %d' % epoch)
      net.train()
      train_loss = 0
      correct = 0
      total = 0
      for batch_idx, (inputs, targets) in enumerate(trainloader):
          if use_cuda:
              inputs, targets = inputs.cuda(), targets.cuda()
          # generate mixed inputs, two one-hot label vectors and mixing coefficient
          inputs, targets_a, targets_b, lam = mixup_data(inputs, targets, args.alpha, use_cuda)
          optimizer.zero_grad()
          outputs = net(inputs)
  
          loss_func = mixup_criterion(targets_a, targets_b, lam)
          loss = loss_func(criterion, outputs)
          loss.backward()
          optimizer.step()
  
          train_loss += loss.item()
          _, predicted = torch.max(outputs.data, 1)
          total += targets.size(0)
          correct += (lam * predicted.eq(targets_a.data).float()).cpu().sum() + ((1 - lam) * predicted.eq(targets_b.data).float()).cpu().sum()
          acc = 100.*float(correct)/float(total)
  
          progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
              % (train_loss/(batch_idx+1), acc, correct, total))
  
      return (train_loss/batch_idx, acc)
  
  def test(epoch):
      global best_acc
      net.eval()
      test_loss = 0
      correct = 0
      total = 0
      with torch.no_grad():
          for batch_idx, (inputs, targets) in enumerate(testloader):
              if use_cuda:
                  inputs, targets = inputs.cuda(), targets.cuda()
              outputs = net(inputs)
              loss = cel(outputs, targets)
  
              test_loss += loss.item()
              _, predicted = torch.max(outputs.data, 1)
              total += targets.size(0)
              correct += predicted.eq(targets.data).cpu().sum()
  
              progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                  % (test_loss/(batch_idx+1), 100.*float(correct)/float(total), correct, total))
  
          # Save checkpoint.
          acc = 100.*float(correct)/float(total)
          if acc > best_acc:
              best_acc = acc
              checkpoint(acc, epoch)
  
      return (test_loss/batch_idx, acc)
  
  def checkpoint(acc, epoch):
      # Save checkpoint.
      print('Saving..')
      state = {
          'net': net,
          'acc': acc,
          'epoch': epoch,
          'rng_state': torch.get_rng_state()
      }
      if not os.path.isdir(os.path.join(args.logdir, 'checkpoint/')):
          os.mkdir(os.path.join(args.logdir, 'checkpoint/'))
      torch.save(state, os.path.join(args.logdir, 'checkpoint/') +
                 args.arch + '_' + args.sess + '_' + str(args.seed) + '.ckpt')
  
  def adjust_learning_rate(optimizer, epoch):
      """decrease the learning rate at 100 and 150 epoch"""
      lr = base_learning_rate
      if epoch <= 9 and lr > 0.1:
          # warm-up training for large minibatch
          lr = 0.1 + (base_learning_rate - 0.1) * epoch / 10.
      if epoch >= 100:
          lr /= 10
      if epoch >= 150:
          lr /= 10
      for param_group in optimizer.param_groups:
          if param_group['initial_lr'] == base_learning_rate:
              param_group['lr'] = lr
          else:
              if epoch <= 9:
                  param_group['lr'] = param_group['initial_lr'] * lr / base_learning_rate
              elif epoch < 100:
                  param_group['lr'] = param_group['initial_lr']
              elif epoch < 150:
                  param_group['lr'] = param_group['initial_lr'] / 10.
              else:
                  param_group['lr'] = param_group['initial_lr'] / 100.
      return lr
  
  if not os.path.exists(logname):
      with open(logname, 'w') as logfile:
          logwriter = csv.writer(logfile, delimiter=',')
          logwriter.writerow(['epoch', 'lr', 'train loss', 'train acc', 'test loss', 'test acc'])
  
  sgdr = CosineAnnealingLR(optimizer, args.n_epoch, eta_min=0, last_epoch=-1)
  
  for epoch in range(start_epoch, args.n_epoch):
      lr = 0.
      if args.sgdr:
          sgdr.step()
          for param_group in optimizer.param_groups:
              lr = param_group['lr']
              break
      else:
          lr = adjust_learning_rate(optimizer, epoch)
      train_loss, train_acc = train(epoch)
      test_loss, test_acc = test(epoch)
      with open(logname, 'a') as logfile:
          logwriter = csv.writer(logfile, delimiter=',')
          logwriter.writerow([epoch, lr, train_loss, train_acc, test_loss, test_acc])
  