autoattack_example.py

#!/usr/bin/env python
# coding: utf-8
import numpy as np
import random
from PIL import Image
import matplotlib.pyplot as plt

import torch#!/usr/bin/env python
# coding: utf-8
import numpy as np
import random
from PIL import Image
import matplotlib.pyplot as plt

import torch
from torch import nn
import torchvision
import torchvision.transforms as T

from autoattack import AutoAttack

# Data
dataset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True)
cifar10_categories = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]

img2tensor = T.Compose([ # to make a tensor in [0, 1]
                       T.ToTensor()])

# Prepare a set of images
NUM_IMG_TO_ATTACK = 100
torch.manual_seed(1)
indices = torch.randperm(len(dataset))[:NUM_IMG_TO_ATTACK]

X = []
y = []
for i in indices:
    im, label = dataset[i]
    im_tensor = img2tensor(im)
    X.append(im_tensor.unsqueeze(0))
    y.append(label)
X = torch.vstack(X)
y = torch.tensor(y)


# load whatever Model: input is normalized -- x = (x - mean) / std
import torchvision.models as models
model = models.resnet18(pretrained=True)


# Model wrapper
class ModelWrapper(nn.Module):
    def __init__(self, model):
        super().__init__()

        self.model = model
        self.data_categories = cifar10_categories
        self.transforms = T.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616])

    def forward(self, x):
        if self.transforms is not None:
            x = self.transforms(x)

        out = self.model(x) # logits of shape (B, num_classes)
        return out

    def predict_class(self, input_tensor, verbose=False):
        out = self.forward(input_tensor)

        scores = torch.nn.functional.softmax(out, dim=-1)[0]

        topk_scores, topk_label = torch.topk(scores, k=3, dim=-1)

        if verbose:
            for i in range(3):
                pred_name = self.data_categories[topk_label[i]]
                print(f"Prediction index {i}: {pred_name:<20}, score: {topk_scores[i].item():.3f}")

        return topk_label[0]

# The input pixel values for this model wrapper have to be in range [0, 1]
device = torch.device('cuda:0')
model_wrapper = ModelWrapper(model).to(device)


# Batch attack
x_advs = []
for eps in [0.0001, 0.0003, 0.0005, 0.001, 0.003, 0.005, 0.01, 8/255]:
    print(f'eps = {eps}')
    adversary = AutoAttack(model_wrapper, norm='Linf', eps=eps, verbose=True, version='standard', device=device)
    adversary.attacks_to_run = ["apgd-t"] # only targeted PGD attack
    x_adv = adversary.run_standard_evaluation(X, y, bs=10)
    x_advs.append(x_adv)

torch.save(x_advs, 'adversarial_examples.pt')
from torch import nn
import torchvision
import torchvision.transforms as T

from autoattack import AutoAttack

# Data
dataset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True)
cifar10_categories = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]

img2tensor = T.Compose([ # to make a tensor in [0, 1]
                       T.ToTensor()])

# Prepare a set of images
NUM_IMG_TO_ATTACK = 100
torch.manual_seed(1)
indices = torch.randperm(len(dataset))[:NUM_IMG_TO_ATTACK]

X = []
y = []
for i in indices:
    im, label = dataset[i]
    im_tensor = img2tensor(im)
    X.append(im_tensor.unsqueeze(0))
    y.append(label)
X = torch.vstack(X)
y = torch.tensor(y)


# load whatever Model: input is normalized -- x = (x - mean) / std
model = WhateverModel(...)


# Model wrapper
class ModelWrapper(nn.Module):
    def __init__(self, model):
        super().__init__()

        self.model = model
        self.data_categories = cifar10_categories
        self.transforms = T.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616])

    def forward(self, x):
        if self.transforms is not None:
            x = self.transforms(x)

        out = self.model(x) # logits of shape (B, num_classes)
        return out

    def predict_class(self, input_tensor, verbose=False):
        out = self.forward(input_tensor)

        scores = torch.nn.functional.softmax(out, dim=-1)[0]

        topk_scores, topk_label = torch.topk(scores, k=3, dim=-1)

        if verbose:
            for i in range(3):
                pred_name = self.data_categories[topk_label[i]]
                print(f"Prediction index {i}: {pred_name:<20}, score: {topk_scores[i].item():.3f}")

        return topk_label[0]

# The input pixel values for this model wrapper have to be in range [0, 1]
model_wrapper = ModelWrapper(model)


# Batch attack
device = torch.device('cuda:0')
x_advs = []
for eps in [0.0001, 0.0003, 0.0005, 0.001, 0.003, 0.005, 0.01, 8/255]:
    print(f'eps = {eps}')
    adversary = AutoAttack(model_wrapper, norm='Linf', eps=eps, verbose=True, version='standard', device=device)
    adversary.attacks_to_run = ["apgd-t"] # only targeted PGD attack
    x_adv = adversary.run_standard_evaluation(X, y, bs=10)
    x_advs.append(x_adv)

torch.save(x_advs, 'adversarial_examples.pt')