first commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 XXX
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,92 @@
+# Adversarial Distributional Training
+
+This repository contains the code for adversarial distributional training (ADT) of our submission: *Adversarial Distributional Training for Robust Deep Learning*, to NeurIPS 2020.
+
+<img src="algos_adt.pdf">
+
+Figure 1: An illustration of three different ADT methods, including (a) ADT<sub>EXP</sub>; (b) ADT<sub>EXP-AM</sub>; (c) ADT<sub>IMP-AM</sub>.
+
+
+
+## Prerequisites
+* Python (3.6.8)
+* Pytorch (1.3.0)
+* torchvision (0.4.1)
+* numpy
+
+## Training
+
+We have proposed three different methods for ADT. The command for each training method is specified below.
+
+### Training ADT<sub>EXP</sub>
+
+```
+python adt_exp.py --model-dir adt-exp --dataset cifar10 (or cifar100/svhn)
+```
+
+### Training ADT<sub>EXP-AM</sub>
+
+```
+python adt_expam.py --model-dir adt-expam --dataset cifar10 (or cifar100/svhn)
+```
+
+### Training ADT<sub>IMP-AM</sub>
+
+```
+python adt_impam.py --model-dir adt-impam --dataset cifar10 (or cifar100/svhn)
+```
+
+The checkpoints will be saved at each model folder.
+
+## Evaluation
+
+### Evaluation under White-box Attacks
+
+- For FGSM attack, run
+
+```
+python evaluate_attacks.py --model-path ${MODEL-PATH} --attack-method FGSM --dataset cifar10 (or cifar100/svhn)
+```
+
+- For PGD attack, run
+
+```
+python evaluate_attacks.py --model-path ${MODEL-PATH} --attack-method PGD --num-steps 20 (or 100) --dataset cifar10 (or cifar100/svhn)
+```
+
+- For MIM attack, run
+
+```
+python evaluate_attacks.py --model-path ${MODEL-PATH} --attack-method MIM --num-steps 20 --dataset cifar10 (or cifar100/svhn)
+```
+
+- For C&W attack, run
+
+```
+python evaluate_attacks.py --model-path ${MODEL-PATH} --attack-method CW --num-steps 30 --dataset cifar10 (or cifar100/svhn)
+```
+
+- For FeaAttack, run
+
+```
+python feature_attack.py --model-path ${MODEL-PATH} --dataset cifar10 (or cifar100/svhn)
+```
+
+
+### Evaluation under Transfer-based Black-box Attacks
+
+First change the `--white-box-attack` argument in `evaluate_attacks.py` to `False`. Then run
+
+```
+python evaluate_attacks.py --source-model-path ${SOURCE-MODEL-PATH} --target-model-path ${TARGET-MODEL-PATH} --attack-method PGD (or MIM)
+```
+
+### Evaluation under SPSA
+
+```
+python spsa.py --model-path ${MODEL-PATH} --samples_per_draw 256 (or 512/1024/2048)
+```
+
+## Pretrained Models
+
+We will release the pretrained models after the review process. It's easy to reproduce the results in our paper, as ADT<sub>EXP-AM</sub> and ADT<sub>IMP-AM</sub> need less than one GPU day to finish training.

adt_exp.py

@@ -0,0 +1,235 @@
+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 Adversarial Distributional 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()