add codes

Author: wzekai99

README.md

@@ -1,2 +1,100 @@
-# Diffusion-Models-Improve-AT
-Code for the paper "Better Diffusion Models Further Improve Adversarial Training"
+# Better Diffusion Models Further Improve Adversarial Training
+
+
+
+## Environment settings and libraries we used in our experiments
+
+This project is tested under the following environment settings:
+- OS: Ubuntu 20.04.3
+- GPU: NVIDIA A100
+- Cuda: 11.1, Cudnn: v8.2
+- Python: 3.9.5
+- PyTorch: 1.8.0
+- Torchvision: 0.9.0
+
+## Acknowledgement
+The codes are modifed based on the [PyTorch implementation](https://github.com/imrahulr/adversarial_robustness_pytorch) of [Rebuffi et al., 2021](https://arxiv.org/abs/2103.01946).
+
+## Requirements
+
+- Install or download [AutoAttack](https://github.com/fra31/auto-attack):
+```
+pip install git+https://github.com/fra31/auto-attack
+```
+
+- Install or download [RandAugment](https://github.com/ildoonet/pytorch-randaugment):
+```
+pip install git+https://github.com/ildoonet/pytorch-randaugment
+```
+
+- Download EDM generated data. Since 20M and 50M data files are too large, we split them into several parts:
+
+| dataset | size | link |
+|---|:---:|:---:|
+| CIFAR-10 | 1M | [npz](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+| CIFAR-10 | 5M | [npz](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+| CIFAR-10 | 10M | [npz](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+| CIFAR-10 | 20M | [part1](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part2](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+| CIFAR-10 | 50M | [part1](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part2](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part3](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part4](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+| CIFAR-100 | 1M | [npz](https://storage.googleapis.com/dm-adversarial-robustness/cifar100_ddpm.npz) |
+| CIFAR-100 | 50M | [part1](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part2](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part3](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) [part4](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_ddpm.npz) |
+
+- Merge 20M and 50M generated data: 
+  
+```
+python merge_data.py
+```
+
+## Training Commands
+
+Run [`train-wa.py`](./train-wa.py) for reproducing the results reported in the papers. For example, train a WideResNet-28-10 model via [TRADES](https://github.com/yaodongyu/TRADES) on CIFAR-10 with the additional generated data provided by EDM ([Karras et al., 2022](https://github.com/NVlabs/edm)):
+
+```python
+python train-wa.py --data-dir 'cifar-data' \
+    --log-dir 'trained_models' \
+    --desc 'WRN28-10Swish_cifar10s_lr0p2_TRADES5_epoch400_bs512_fraction0p7_ls0p1' \
+    --data cifar10s \
+    --batch-size 512 \
+    --model wrn-28-10-swish \
+    --num-adv-epochs 400 \
+    --lr 0.2 \
+    --beta 5.0 \
+    --unsup-fraction 0.7 \
+    --aux-data-filename <path_to_additional_data> \
+    --ls 0.1
+```
+
+
+
+## Downloading models
+
+We provide checkpoints which  Download a model from links listed in the following table. Clean and robust accuracies are measured on the full test set. The robust accuracy is measured using [AutoAttack](https://github.com/fra31/auto-attack).
+
+| dataset | norm | radius | architecture | clean | robust | link |
+|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| CIFAR-10 | &#8467;<sub>&infin;</sub> | 8 / 255 | WRN-28-10 | 92.44% | 67.31% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+| CIFAR-10 | &#8467;<sub>&infin;</sub> | 8 / 255 | WRN-70-16 | 93.25% | 70.69% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn70-16_with.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+| CIFAR-10 | &#8467;<sub>2</sub> | 128 / 255 | WRN-28-10 | 95.16% | 83.63% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_l2_wrn70-16_with.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+| CIFAR-10 | &#8467;<sub>2</sub> | 128 / 255 | WRN-70-16 | 95.54% | 84.86% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_l2_wrn70-16_without.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+| CIFAR-100 | &#8467;<sub>&infin;</sub> | 8 / 255 | WRN-28-10 | 72.58% | 38.83% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar100_linf_wrn70-16_with.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+| CIFAR-100 | &#8467;<sub>&infin;</sub> | 8 / 255 | WRN-70-16 | 75.22% | 42.67% | [checkpoint](https://storage.googleapis.com/dm-adversarial-robustness/cifar100_linf_wrn70-16_without.pt) [argtxt](https://storage.googleapis.com/dm-adversarial-robustness/cifar10_linf_wrn28-10_with.pt)
+
+- **Downloading `checkpoint` to `trained_models/mymodel/weights-best.pt`**
+- **Downloading `argtxt` to `trained_models/mymodel/args.txt`**
+
+## Evaluation Commands
+The trained models can be evaluated by running [`eval-aa.py`](./eval-aa.py) which uses [AutoAttack](https://github.com/fra31/auto-attack) for evaluating the robust accuracy. Run the command:
+
+```python
+python eval-aa.py --data-dir 'cifar-data' \
+    --log-dir 'trained_models' \
+    --desc mymodel
+```
+
+To evaluate the model on last epoch under AutoAttack, run the command: 
+
+```python
+python eval-last-aa.py --data-dir 'cifar-data' \
+    --log-dir 'trained_models' \
+    --desc mymodel
+```

core/__init__.py

@@ -0,0 +1 @@
+

core/attacks/__init__.py

@@ -0,0 +1,65 @@
+from .base import Attack
+
+from .apgd import LinfAPGDAttack
+from .apgd import L2APGDAttack
+
+from .fgsm import FGMAttack
+from .fgsm import FGSMAttack
+from .fgsm import L2FastGradientAttack
+from .fgsm import LinfFastGradientAttack
+
+from .pgd import PGDAttack
+from .pgd import L2PGDAttack
+from .pgd import LinfPGDAttack
+
+from .deepfool import DeepFoolAttack
+from .deepfool import LinfDeepFoolAttack
+from .deepfool import L2DeepFoolAttack
+
+from .utils import CWLoss
+
+
+ATTACKS = ['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd']
+
+
+def create_attack(model, criterion, attack_type, attack_eps, attack_iter, attack_step, rand_init_type='uniform', 
+                  clip_min=0., clip_max=1.):
+    """
+    Initialize adversary.
+    Arguments:
+        model (nn.Module): forward pass function.
+        criterion (nn.Module): loss function.
+        attack_type (str): name of the attack.
+        attack_eps (float): attack radius.
+        attack_iter (int): number of attack iterations.
+        attack_step (float): step size for the attack.
+        rand_init_type (str): random initialization type for PGD (default: uniform).
+        clip_min (float): mininum value per input dimension.
+        clip_max (float): maximum value per input dimension.
+   Returns:
+       Attack
+   """
+    
+    if attack_type == 'fgsm':
+        attack = FGSMAttack(model, criterion, eps=attack_eps, clip_min=clip_min, clip_max=clip_max)
+    elif attack_type == 'fgm':
+        attack = FGMAttack(model, criterion, eps=attack_eps, clip_min=clip_min, clip_max=clip_max)
+    elif attack_type == 'linf-pgd':
+        attack = LinfPGDAttack(model, criterion, eps=attack_eps, nb_iter=attack_iter, eps_iter=attack_step,
+                               rand_init_type=rand_init_type, clip_min=clip_min, clip_max=clip_max)
+    elif attack_type == 'l2-pgd':
+        attack = L2PGDAttack(model, criterion, eps=attack_eps, nb_iter=attack_iter, eps_iter=attack_step, 
+                             rand_init_type=rand_init_type, clip_min=clip_min, clip_max=clip_max)
+    elif attack_type == 'linf-df':
+        attack = LinfDeepFoolAttack(model, overshoot=0.02, nb_iter=attack_iter, search_iter=0, clip_min=clip_min, 
+                                    clip_max=clip_max)
+    elif attack_type == 'l2-df':
+        attack = L2DeepFoolAttack(model, overshoot=0.02, nb_iter=attack_iter, search_iter=0, clip_min=clip_min, 
+                                  clip_max=clip_max)
+    elif attack_type == 'linf-apgd':
+        attack = LinfAPGDAttack(model, criterion, n_restarts=2, eps=attack_eps, nb_iter=attack_iter)
+    elif attack_type == 'l2-apgd':
+        attack = L2APGDAttack(model, criterion, n_restarts=2, eps=attack_eps, nb_iter=attack_iter)
+    else:
+        raise NotImplementedError('{} is not yet implemented!'.format(attack_type))
+    return attack

core/attacks/apgd.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+import torch
+from autoattack.autopgd_base import APGDAttack
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class APGD():
+    """
+    APGD attack (from AutoAttack) (Croce et al, 2020).
+    The attack performs nb_iter steps of adaptive size, while always staying within eps from the initial point.
+    Arguments:
+        predict (nn.Module): forward pass function.
+        loss_fn (str): loss function - ce or dlr.
+        n_restarts (int): number of random restarts.
+        eps (float): maximum distortion.
+        nb_iter (int): number of iterations.
+        ord (int): (optional) the order of maximum distortion (inf or 2).
+    """
+    def __init__(self, predict, loss_fn='ce', n_restarts=2, eps=0.3, nb_iter=40, ord=np.inf, seed=1):
+        assert loss_fn in ['ce', 'dlr'], 'Only loss_fn=ce or loss_fn=dlr are supported!'
+        assert ord in [2, np.inf], 'Only ord=inf or ord=2 are supported!'
+        
+        norm = 'Linf' if ord == np.inf else 'L2'
+        self.apgd = APGDAttack(predict, n_restarts=n_restarts, n_iter=nb_iter, verbose=False, eps=eps, norm=norm, 
+                               eot_iter=1, rho=.75, seed=seed, device=device)
+        self.apgd.loss = loss_fn
+
+    def perturb(self, x, y):
+        x_adv = self.apgd.perturb(x, y)[1]
+        r_adv = x_adv - x
+        return x_adv, r_adv
+
+    
+class LinfAPGDAttack(APGD):
+    """
+    APGD attack (from AutoAttack) with order=Linf.
+    The attack performs nb_iter steps of adaptive size, while always staying within eps from the initial point.
+    Arguments:
+        predict (nn.Module): forward pass function.
+        loss_fn (str): loss function - ce or dlr.
+        n_restarts (int): number of random restarts.
+        eps (float): maximum distortion.
+        nb_iter (int): number of iterations.
+    """
+    
+    def __init__(self, predict, loss_fn='ce', n_restarts=2, eps=0.3, nb_iter=40, seed=1):
+        ord = np.inf
+        super(L2APGDAttack, self).__init__(
+            predict=predict, loss_fn=loss_fn, n_restarts=n_restarts, eps=eps, nb_iter=nb_iter, ord=ord, seed=seed)
+
+
+class L2APGDAttack(APGD):
+    """
+    APGD attack (from AutoAttack) with order=L2.
+    The attack performs nb_iter steps of adaptive size, while always staying within eps from the initial point.
+    Arguments:
+        predict (nn.Module): forward pass function.
+        loss_fn (str): loss function - ce or dlr.
+        n_restarts (int): number of random restarts.
+        eps (float): maximum distortion.
+        nb_iter (int): number of iterations.
+    """
+    
+    def __init__(self, predict, loss_fn='ce', n_restarts=2, eps=0.3, nb_iter=40, seed=1):
+        ord = 2
+        super(L2APGDAttack, self).__init__(
+            predict=predict, loss_fn=loss_fn, n_restarts=n_restarts, eps=eps, nb_iter=nb_iter, ord=ord, seed=seed)

core/attacks/base.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+
+from .utils import replicate_input
+
+
+class Attack(object):
+    """
+    Abstract base class for all attack classes.
+    Arguments:
+        predict (nn.Module): forward pass function.
+        loss_fn (nn.Module): loss function.
+        clip_min (float): mininum value per input dimension.
+        clip_max (float): maximum value per input dimension.
+    """
+
+    def __init__(self, predict, loss_fn, clip_min, clip_max):
+        self.predict = predict
+        self.loss_fn = loss_fn
+        self.clip_min = clip_min
+        self.clip_max = clip_max
+
+    def perturb(self, x, **kwargs):
+        """
+        Virtual method for generating the adversarial examples.
+        Arguments:
+            x (torch.Tensor): the model's input tensor.
+            **kwargs: optional parameters used by child classes.
+        Returns: 
+            adversarial examples.
+        """
+        error = "Sub-classes must implement perturb."
+        raise NotImplementedError(error)
+
+    def __call__(self, *args, **kwargs):
+        return self.perturb(*args, **kwargs)
+
+
+class LabelMixin(object):
+    def _get_predicted_label(self, x):
+        """
+        Compute predicted labels given x. Used to prevent label leaking during adversarial training.
+        Arguments:
+            x (torch.Tensor): the model's input tensor.
+        Returns:
+            torch.Tensor containing predicted labels.
+        """
+        with torch.no_grad():
+            outputs = self.predict(x)
+        _, y = torch.max(outputs, dim=1)
+        return y
+
+    def _verify_and_process_inputs(self, x, y):
+        if self.targeted:
+            assert y is not None
+
+        if not self.targeted:
+            if y is None:
+                y = self._get_predicted_label(x)
+
+        x = replicate_input(x)
+        y = replicate_input(y)
+        return x, y