Compare commits
No commits in common. "vision_only" and "branch1" have entirely different histories.
vision_onl
...
branch1
132
auto_cifar10.py
132
auto_cifar10.py
|
|
@ -1,132 +0,0 @@
|
||||||
import warnings
|
|
||||||
|
|
||||||
import torchvision.datasets
|
|
||||||
|
|
||||||
warnings.filterwarnings('ignore')
|
|
||||||
|
|
||||||
from PIL import Image
|
|
||||||
import torch
|
|
||||||
import timm
|
|
||||||
import requests
|
|
||||||
import numpy as np
|
|
||||||
import torchvision.transforms as transforms
|
|
||||||
from torch import nn
|
|
||||||
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
|
|
||||||
from torch.utils.data import Dataset, DataLoader
|
|
||||||
import copy
|
|
||||||
|
|
||||||
from art.estimators.classification import PyTorchClassifier
|
|
||||||
from art.data_generators import PyTorchDataGenerator
|
|
||||||
from art.utils import load_cifar10
|
|
||||||
from art.attacks.evasion import ProjectedGradientDescent ,AutoProjectedGradientDescent
|
|
||||||
from art.defences.trainer import AdversarialTrainer
|
|
||||||
|
|
||||||
model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k", pretrained=False)
|
|
||||||
model.head = nn.Linear(model.head.in_features, 10)
|
|
||||||
state_dict = torch.load('/home/leewlving/.cache/torch/hub/checkpoints/vit_base_patch16_224_in21k_ft_cifar10.pth')
|
|
||||||
model.load_state_dict(state_dict)
|
|
||||||
# model.load_state_dict(
|
|
||||||
# torch.hub.load_state_dict_from_url(
|
|
||||||
# "https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar10/resolve/main/pytorch_model.bin",
|
|
||||||
# map_location="cuda",
|
|
||||||
# file_name="vit_base_patch16_224_in21k_ft_cifar10.pth",
|
|
||||||
# )
|
|
||||||
# )
|
|
||||||
model.eval()
|
|
||||||
|
|
||||||
DEFAULT_MEAN = (0.485, 0.456, 0.406)
|
|
||||||
DEFAULT_STD = (0.229, 0.224, 0.225)
|
|
||||||
|
|
||||||
transform = transforms.Compose([
|
|
||||||
transforms.Resize(256, interpolation=3),
|
|
||||||
transforms.CenterCrop(224),
|
|
||||||
transforms.ToTensor(),
|
|
||||||
transforms.Normalize(DEFAULT_MEAN, DEFAULT_STD),
|
|
||||||
])
|
|
||||||
|
|
||||||
|
|
||||||
class CIFAR10_dataset(Dataset):
|
|
||||||
def __init__(self, data, targets, transform=None):
|
|
||||||
self.data = data
|
|
||||||
self.targets = torch.LongTensor(targets)
|
|
||||||
self.transform = transform
|
|
||||||
|
|
||||||
def __getitem__(self, index):
|
|
||||||
x = Image.fromarray(((self.data[index] * 255).round()).astype(np.uint8).transpose(1, 2, 0))
|
|
||||||
x = self.transform(x)
|
|
||||||
y = self.targets[index]
|
|
||||||
return x, y
|
|
||||||
|
|
||||||
def __len__(self):
|
|
||||||
return len(self.data)
|
|
||||||
|
|
||||||
|
|
||||||
# (x_train, y_train), (x_test, y_test), min_pixel_value, max_pixel_value = load_cifar10()
|
|
||||||
# print(max_pixel_value)
|
|
||||||
# x_train = x_train.transpose(0, 3, 1, 2).astype("float32")
|
|
||||||
# x_test = x_test.transpose(0, 3, 1, 2).astype("float32")
|
|
||||||
train_dataset = torchvision.datasets.SVHN(root='./svhn',split='train',download=True,transform=transform)
|
|
||||||
test_dataset= torchvision.datasets.SVHN(root='./svhn',split='test',download=True,transform=transform)
|
|
||||||
# dataset = CIFAR10_dataset(x_train, y_train, transform=transform)
|
|
||||||
dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
|
|
||||||
test_dataloader =DataLoader(test_dataset, batch_size=64, shuffle=False)
|
|
||||||
|
|
||||||
opt = torch.optim.Adam(model.parameters(), lr=0.01)
|
|
||||||
|
|
||||||
|
|
||||||
criterion = nn.CrossEntropyLoss()
|
|
||||||
|
|
||||||
classifier = PyTorchClassifier(
|
|
||||||
model=model,
|
|
||||||
clip_values=(0.0, 1.0),
|
|
||||||
loss=criterion,
|
|
||||||
optimizer=opt,
|
|
||||||
input_shape=(3, 224, 224),
|
|
||||||
nb_classes=10,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
attack= AutoProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=10,
|
|
||||||
targeted=False,
|
|
||||||
batch_size=64,
|
|
||||||
verbose=False
|
|
||||||
)
|
|
||||||
# attack = ProjectedGradientDescent(
|
|
||||||
# classifier,
|
|
||||||
# norm=np.inf,
|
|
||||||
# eps=8.0 / 255.0,
|
|
||||||
# eps_step=2.0 / 255.0,
|
|
||||||
# max_iter=10,
|
|
||||||
# targeted=False,
|
|
||||||
# num_random_init=1,
|
|
||||||
# batch_size=64,
|
|
||||||
# verbose=False,
|
|
||||||
# )
|
|
||||||
|
|
||||||
trainer = AdversarialTrainer(
|
|
||||||
classifier, attack
|
|
||||||
)
|
|
||||||
art_datagen = PyTorchDataGenerator(iterator=dataloader, size=len(train_dataset), batch_size=64)
|
|
||||||
|
|
||||||
trainer.fit_generator(art_datagen, nb_epochs=1)
|
|
||||||
|
|
||||||
# for i, data in enumerate(test_dataloader):
|
|
||||||
# x, y = data
|
|
||||||
# x = x.numpy()
|
|
||||||
# y = y.numpy()
|
|
||||||
# # print(x.shape)
|
|
||||||
# # print(y.shape)
|
|
||||||
# x_test_pred = np.argmax(classifier.predict(x), axis=1)
|
|
||||||
# print(
|
|
||||||
# "Accuracy on benign test samples after adversarial training: %.2f%%"
|
|
||||||
# % (np.sum(x_test_pred == np.argmax(y, axis=1)) / x.shape[0] * 100)
|
|
||||||
# )export https_proxy=http://127.0.0.1:7897 http_proxy=http://127.0.0.1:7897 all_proxy=socks5://127.0.0.1:7897
|
|
||||||
|
|
||||||
# trainer.classifier.save('AT-cifar10.pth')
|
|
||||||
torch.save(trainer.classifier.model.state_dict(), 'AT-svhn.pth')
|
|
||||||
|
|
||||||
192
cifar10.py
192
cifar10.py
|
|
@ -1,192 +0,0 @@
|
||||||
from PIL import Image
|
|
||||||
import numpy as np
|
|
||||||
import timm
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
import torch.nn.functional as F
|
|
||||||
import torchvision.transforms as transforms
|
|
||||||
from torch.utils.data import Dataset, DataLoader
|
|
||||||
from torch.optim.lr_scheduler import MultiStepLR, StepLR
|
|
||||||
|
|
||||||
from art.estimators.classification import PyTorchClassifier
|
|
||||||
from art.data_generators import PyTorchDataGenerator
|
|
||||||
from art.defences.trainer import AdversarialTrainer
|
|
||||||
from art.attacks.evasion import ProjectedGradientDescent
|
|
||||||
from datasets import load_dataset
|
|
||||||
from torchvision.transforms import (CenterCrop,
|
|
||||||
Compose,
|
|
||||||
Normalize,
|
|
||||||
RandomHorizontalFlip,
|
|
||||||
RandomResizedCrop,
|
|
||||||
Resize,
|
|
||||||
ToTensor)
|
|
||||||
from tensorflow.keras.utils import to_categorical
|
|
||||||
from transformers import ViTImageProcessor
|
|
||||||
|
|
||||||
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
|
|
||||||
IMAGENET_DEFAULT_MEAN = processor.image_mean
|
|
||||||
IMAGENET_DEFAULT_STD = processor.image_std
|
|
||||||
|
|
||||||
size = processor.size["height"]
|
|
||||||
|
|
||||||
|
|
||||||
"""
|
|
||||||
For this example we choose the ResNet18 model as used in the paper (https://proceedings.mlr.press/v97/zhang19p.html)
|
|
||||||
The code for the model architecture has been adopted from
|
|
||||||
https://github.com/yaodongyu/TRADES/blob/master/models/resnet.py
|
|
||||||
"""
|
|
||||||
|
|
||||||
|
|
||||||
model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k", pretrained=False)
|
|
||||||
model.head = nn.Linear(model.head.in_features, 10)
|
|
||||||
model.load_state_dict(
|
|
||||||
torch.hub.load_state_dict_from_url(
|
|
||||||
"https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar10/resolve/main/pytorch_model.bin",
|
|
||||||
map_location="cuda",
|
|
||||||
file_name="vit_base_patch16_224_in21k_ft_cifar10.pth",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# Step 1: Load the CIFAR10 dataset
|
|
||||||
train_ds, test_ds = load_dataset('cifar10', split=['train[:5000]', 'test[:2000]'])
|
|
||||||
splits = train_ds.train_test_split(test_size=0.1)
|
|
||||||
train_ds = splits['train']
|
|
||||||
val_ds = splits['test']
|
|
||||||
|
|
||||||
train_size=len(train_ds)
|
|
||||||
test_size=len(test_ds)
|
|
||||||
|
|
||||||
normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
|
|
||||||
_train_transforms = Compose(
|
|
||||||
[
|
|
||||||
RandomResizedCrop(size),
|
|
||||||
RandomHorizontalFlip(),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
_val_transforms = Compose(
|
|
||||||
[
|
|
||||||
Resize(size),
|
|
||||||
CenterCrop(size),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
def train_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['img']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
def val_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['img']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
train_ds.set_transform(train_transforms)
|
|
||||||
val_ds.set_transform(val_transforms)
|
|
||||||
test_ds.set_transform(val_transforms)
|
|
||||||
|
|
||||||
|
|
||||||
def collate_fn(examples):
|
|
||||||
pixel_values = torch.stack([example["pixel_values"] for example in examples])
|
|
||||||
labels = torch.tensor([example["label"] for example in examples])
|
|
||||||
return pixel_values,labels
|
|
||||||
|
|
||||||
train_batch_size = 32
|
|
||||||
eval_batch_size = 32
|
|
||||||
|
|
||||||
def dataset2np(dataset):
|
|
||||||
X = []
|
|
||||||
Y = []
|
|
||||||
for i in range(int(2000)):
|
|
||||||
x,y = dataset[i]["pixel_values"], dataset[i]["label"]
|
|
||||||
y=to_categorical(y,num_classes=10)
|
|
||||||
X.append(x.detach().numpy())
|
|
||||||
Y.append(y)
|
|
||||||
X = np.array(X).astype("float32")
|
|
||||||
Y = np.array(Y).astype("float32")
|
|
||||||
return X,Y
|
|
||||||
|
|
||||||
train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
|
|
||||||
val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
x_test, y_test=dataset2np(test_ds)
|
|
||||||
|
|
||||||
|
|
||||||
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
|
|
||||||
lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
|
|
||||||
|
|
||||||
criterion = nn.CrossEntropyLoss()
|
|
||||||
|
|
||||||
# Step 3: Create the ART classifier
|
|
||||||
|
|
||||||
classifier = PyTorchClassifier(
|
|
||||||
model=model,
|
|
||||||
clip_values=(0.0, 1.0),
|
|
||||||
loss=criterion,
|
|
||||||
optimizer=opt,
|
|
||||||
input_shape=(3, size, size),
|
|
||||||
nb_classes=10,
|
|
||||||
)
|
|
||||||
|
|
||||||
attack = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=10,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on clean samples before adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
|
|
||||||
trainer = AdversarialTrainer(
|
|
||||||
classifier, attack
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build a Keras image augmentation object and wrap it in ART
|
|
||||||
art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
|
|
||||||
|
|
||||||
# Step 5: fit the trainer
|
|
||||||
trainer.fit_generator(art_datagen, nb_epochs=50)
|
|
||||||
|
|
||||||
|
|
||||||
x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on benign test samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
attack_test = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=20,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
x_test_attack = attack_test.generate(x_test, y=y_test)
|
|
||||||
x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
torch.save(trainer.classifier.model.state_dict(), 'cifar10_pgd.pth')
|
|
||||||
print(
|
|
||||||
"Save the AT model! "
|
|
||||||
)
|
|
||||||
182
cifar100.py
182
cifar100.py
|
|
@ -1,182 +0,0 @@
|
||||||
from PIL import Image
|
|
||||||
import numpy as np
|
|
||||||
import timm
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
import torch.nn.functional as F
|
|
||||||
import torchvision.transforms as transforms
|
|
||||||
from torch.utils.data import Dataset, DataLoader
|
|
||||||
from torch.optim.lr_scheduler import MultiStepLR, StepLR
|
|
||||||
|
|
||||||
from art.estimators.classification import PyTorchClassifier
|
|
||||||
from art.data_generators import PyTorchDataGenerator
|
|
||||||
from art.defences.trainer import AdversarialTrainer
|
|
||||||
from art.attacks.evasion import ProjectedGradientDescent
|
|
||||||
from datasets import load_dataset
|
|
||||||
from torchvision.transforms import (CenterCrop,
|
|
||||||
Compose,
|
|
||||||
Normalize,
|
|
||||||
RandomHorizontalFlip,
|
|
||||||
RandomResizedCrop,
|
|
||||||
Resize,
|
|
||||||
ToTensor)
|
|
||||||
from tensorflow.keras.utils import to_categorical
|
|
||||||
from transformers import ViTImageProcessor
|
|
||||||
|
|
||||||
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
|
|
||||||
IMAGENET_DEFAULT_MEAN = processor.image_mean
|
|
||||||
IMAGENET_DEFAULT_STD = processor.image_std
|
|
||||||
|
|
||||||
size = processor.size["height"]
|
|
||||||
|
|
||||||
|
|
||||||
model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
|
|
||||||
pretrained=False)
|
|
||||||
model.head = nn.Linear(model.head.in_features, 100)
|
|
||||||
model.load_state_dict(
|
|
||||||
torch.hub.load_state_dict_from_url(
|
|
||||||
"https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_cifar100/resolve/main/pytorch_model.bin",
|
|
||||||
map_location="cuda",
|
|
||||||
file_name="vit_base_patch16_224_in21k_ft_cifar100.pth",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
train_ds = load_dataset("uoft-cs/cifar100",split='train')
|
|
||||||
test_ds = load_dataset("uoft-cs/cifar100",split='test')
|
|
||||||
splits = train_ds.train_test_split(test_size=0.1)
|
|
||||||
train_ds = splits['train']
|
|
||||||
val_ds = splits['test']
|
|
||||||
|
|
||||||
train_size=len(train_ds)
|
|
||||||
test_size=len(test_ds)
|
|
||||||
|
|
||||||
normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
|
|
||||||
_train_transforms = Compose(
|
|
||||||
[
|
|
||||||
RandomResizedCrop(size),
|
|
||||||
RandomHorizontalFlip(),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
_val_transforms = Compose(
|
|
||||||
[
|
|
||||||
Resize(size),
|
|
||||||
CenterCrop(size),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
def train_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['img']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
def val_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['img']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
train_ds.set_transform(train_transforms)
|
|
||||||
val_ds.set_transform(val_transforms)
|
|
||||||
test_ds.set_transform(val_transforms)
|
|
||||||
|
|
||||||
def collate_fn(examples):
|
|
||||||
pixel_values = torch.stack([example["pixel_values"] for example in examples])
|
|
||||||
labels = torch.tensor([example["fine_label"] for example in examples])
|
|
||||||
return pixel_values,labels
|
|
||||||
|
|
||||||
train_batch_size = 64
|
|
||||||
eval_batch_size = 64
|
|
||||||
|
|
||||||
def dataset2np(dataset):
|
|
||||||
X = []
|
|
||||||
Y = []
|
|
||||||
for i in range(int(2000)):
|
|
||||||
x,y = dataset[i]["pixel_values"], dataset[i]["fine_label"]
|
|
||||||
y=to_categorical(y,num_classes=100)
|
|
||||||
X.append(x.detach().numpy())
|
|
||||||
Y.append(y)
|
|
||||||
X = np.array(X).astype("float32")
|
|
||||||
Y = np.array(Y).astype("float32")
|
|
||||||
return X,Y
|
|
||||||
|
|
||||||
train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
|
|
||||||
val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
x_test, y_test=dataset2np(test_ds)
|
|
||||||
|
|
||||||
|
|
||||||
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
|
|
||||||
lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
|
|
||||||
|
|
||||||
criterion = nn.CrossEntropyLoss()
|
|
||||||
|
|
||||||
# Step 3: Create the ART classifier
|
|
||||||
|
|
||||||
classifier = PyTorchClassifier(
|
|
||||||
model=model,
|
|
||||||
clip_values=(0.0, 1.0),
|
|
||||||
loss=criterion,
|
|
||||||
optimizer=opt,
|
|
||||||
input_shape=(3, size, size),
|
|
||||||
nb_classes=100,
|
|
||||||
)
|
|
||||||
|
|
||||||
attack = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=10,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on clean samples before adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
|
|
||||||
trainer = AdversarialTrainer(
|
|
||||||
classifier, attack
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build a Keras image augmentation object and wrap it in ART
|
|
||||||
art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
|
|
||||||
|
|
||||||
# Step 5: fit the trainer
|
|
||||||
trainer.fit_generator(art_datagen, nb_epochs=50)
|
|
||||||
|
|
||||||
|
|
||||||
x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on benign test samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
attack_test = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=20,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
x_test_attack = attack_test.generate(x_test, y=y_test)
|
|
||||||
x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
torch.save(trainer.classifier.model.state_dict(), 'cifar100-pgd.pth')
|
|
||||||
print(
|
|
||||||
"Save the AT model! "
|
|
||||||
)
|
|
||||||
92
mlp_mix.py
92
mlp_mix.py
|
|
@ -1,92 +0,0 @@
|
||||||
import typing as tp
|
|
||||||
|
|
||||||
import jax
|
|
||||||
import jax.numpy as jnp
|
|
||||||
from einops import rearrange
|
|
||||||
from functools import partial
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
from flax import nnx
|
|
||||||
|
|
||||||
class FeedForward(nnx.Module):
|
|
||||||
def __init__(self, dim, hidden_dim, dropout , rngs: nnx.Rngs):
|
|
||||||
self.net=nnx.Sequential(
|
|
||||||
nnx.Linear(dim, hidden_dim , rngs=rngs),
|
|
||||||
partial(nnx.gelu),
|
|
||||||
nnx.Dropout(dropout , rngs=rngs),
|
|
||||||
nnx.Linear(hidden_dim, dim , rngs=rngs),
|
|
||||||
nnx.Dropout(dropout , rngs=rngs)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def __call__(self, x):
|
|
||||||
return self.net(x)
|
|
||||||
|
|
||||||
|
|
||||||
class MixerBlock(nnx.Module):
|
|
||||||
|
|
||||||
def __init__(self, dim, num_patch, token_dim, channel_dim, dropout , rngs: nnx.Rngs):
|
|
||||||
super().__init__()
|
|
||||||
self.ln1=nnx.LayerNorm(dim, rngs=rngs)
|
|
||||||
self.ffn1=FeedForward(num_patch,token_dim,dropout,rngs=rngs)
|
|
||||||
|
|
||||||
self.ln2=nnx.LayerNorm(dim, rngs=rngs)
|
|
||||||
self.ffn2=FeedForward(dim, channel_dim, dropout, rngs=rngs)
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def __call__(self, x):
|
|
||||||
# print(x.shape)
|
|
||||||
x = x + self.ffn1(self.ln1(x))
|
|
||||||
|
|
||||||
x = x + self.ffn2(self.ln2(x))
|
|
||||||
|
|
||||||
return x
|
|
||||||
|
|
||||||
class MLPMixer(nnx.Module):
|
|
||||||
|
|
||||||
def __init__(self, in_channels, dim, num_classes, patch_size,dropout, image_size, depth, token_dim, channel_dim, rngs: nnx.Rngs):
|
|
||||||
super().__init__()
|
|
||||||
|
|
||||||
assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
|
|
||||||
self.num_patch = (image_size// patch_size) ** 2
|
|
||||||
|
|
||||||
self.to_patch_embedding = nnx.Sequential(
|
|
||||||
nnx.Conv(in_channels, dim, kernel_size=(patch_size, patch_size), rngs=rngs),
|
|
||||||
)
|
|
||||||
self.mixer_blocks=[]
|
|
||||||
|
|
||||||
for _ in range(depth):
|
|
||||||
self.mixer_blocks.append(MixerBlock(dim, self.num_patch, token_dim, channel_dim,dropout, rngs=rngs))
|
|
||||||
|
|
||||||
self.layer_norm = nnx.LayerNorm(dim, rngs=rngs)
|
|
||||||
|
|
||||||
self.mlp_head = nnx.Sequential(
|
|
||||||
nnx.Linear(dim, num_classes, rngs=rngs)
|
|
||||||
)
|
|
||||||
|
|
||||||
def __call__(self, x):
|
|
||||||
|
|
||||||
|
|
||||||
x = self.to_patch_embedding(x)
|
|
||||||
|
|
||||||
for mixer_block in self.mixer_blocks:
|
|
||||||
x = mixer_block(x)
|
|
||||||
|
|
||||||
x = self.layer_norm(x)
|
|
||||||
|
|
||||||
x = jnp.mean(x, axis=1)
|
|
||||||
|
|
||||||
return self.mlp_head(x)
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
img = jnp.ones([1, 3, 224, 224])
|
|
||||||
|
|
||||||
model = MLPMixer(in_channels=3, image_size=224, patch_size=16,dropout=0.2, num_classes=1000,
|
|
||||||
dim=512, depth=8, token_dim=256, channel_dim=2048,rngs=nnx.Rngs(0))
|
|
||||||
# nnx.display(model)
|
|
||||||
out_img = model(jnp.ones((1, 224, 224,3)))
|
|
||||||
|
|
||||||
print("Shape of out :", out_img.shape) # [B, in_channels, image_size, image_size]
|
|
||||||
182
svhn_at.py
182
svhn_at.py
|
|
@ -1,182 +0,0 @@
|
||||||
from PIL import Image
|
|
||||||
import numpy as np
|
|
||||||
import timm
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
import torch.nn.functional as F
|
|
||||||
import torchvision.transforms as transforms
|
|
||||||
from torch.utils.data import Dataset, DataLoader
|
|
||||||
from torch.optim.lr_scheduler import MultiStepLR, StepLR
|
|
||||||
|
|
||||||
from art.estimators.classification import PyTorchClassifier
|
|
||||||
from art.data_generators import PyTorchDataGenerator
|
|
||||||
from art.defences.trainer import AdversarialTrainer
|
|
||||||
from art.attacks.evasion import ProjectedGradientDescent
|
|
||||||
from datasets import load_dataset
|
|
||||||
from torchvision.transforms import (CenterCrop,
|
|
||||||
Compose,
|
|
||||||
Normalize,
|
|
||||||
RandomHorizontalFlip,
|
|
||||||
RandomResizedCrop,
|
|
||||||
Resize,
|
|
||||||
ToTensor)
|
|
||||||
from tensorflow.keras.utils import to_categorical
|
|
||||||
from transformers import ViTImageProcessor
|
|
||||||
|
|
||||||
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
|
|
||||||
IMAGENET_DEFAULT_MEAN = processor.image_mean
|
|
||||||
IMAGENET_DEFAULT_STD = processor.image_std
|
|
||||||
|
|
||||||
size = processor.size["height"]
|
|
||||||
|
|
||||||
|
|
||||||
model = timm.create_model("timm/vit_base_patch16_224.orig_in21k_ft_in1k",
|
|
||||||
pretrained=False)
|
|
||||||
model.head = nn.Linear(model.head.in_features, 10)
|
|
||||||
model.load_state_dict(
|
|
||||||
torch.hub.load_state_dict_from_url(
|
|
||||||
"https://huggingface.co/edadaltocg/vit_base_patch16_224_in21k_ft_svhn/resolve/main/pytorch_model.bin",
|
|
||||||
map_location="cuda",
|
|
||||||
file_name="vit_base_patch16_224_in21k_ft_svhn.pth",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
train_ds = load_dataset('ufldl-stanford/svhn', "cropped_digits", split="train")
|
|
||||||
test_ds = load_dataset('ufldl-stanford/svhn', "cropped_digits", split="test")
|
|
||||||
splits = train_ds.train_test_split(test_size=0.1)
|
|
||||||
train_ds = splits['train']
|
|
||||||
val_ds = splits['test']
|
|
||||||
|
|
||||||
train_size=len(train_ds)
|
|
||||||
test_size=len(test_ds)
|
|
||||||
|
|
||||||
normalize = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
|
|
||||||
_train_transforms = Compose(
|
|
||||||
[
|
|
||||||
RandomResizedCrop(size),
|
|
||||||
RandomHorizontalFlip(),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
_val_transforms = Compose(
|
|
||||||
[
|
|
||||||
Resize(size),
|
|
||||||
CenterCrop(size),
|
|
||||||
ToTensor(),
|
|
||||||
normalize,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
def train_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['image']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
def val_transforms(examples):
|
|
||||||
examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['image']]
|
|
||||||
return examples
|
|
||||||
|
|
||||||
train_ds.set_transform(train_transforms)
|
|
||||||
val_ds.set_transform(val_transforms)
|
|
||||||
test_ds.set_transform(val_transforms)
|
|
||||||
|
|
||||||
def collate_fn(examples):
|
|
||||||
pixel_values = torch.stack([example["pixel_values"] for example in examples])
|
|
||||||
labels = torch.tensor([example["label"] for example in examples])
|
|
||||||
return pixel_values,labels
|
|
||||||
|
|
||||||
train_batch_size = 32
|
|
||||||
eval_batch_size = 32
|
|
||||||
|
|
||||||
def dataset2np(dataset):
|
|
||||||
X = []
|
|
||||||
Y = []
|
|
||||||
for i in range(int(2000)):
|
|
||||||
x,y = dataset[i]["pixel_values"], dataset[i]["label"]
|
|
||||||
y=to_categorical(y,num_classes=10)
|
|
||||||
X.append(x.detach().numpy())
|
|
||||||
Y.append(y)
|
|
||||||
X = np.array(X).astype("float32")
|
|
||||||
Y = np.array(Y).astype("float32")
|
|
||||||
return X,Y
|
|
||||||
|
|
||||||
train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
|
|
||||||
val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)
|
|
||||||
x_test, y_test=dataset2np(test_ds)
|
|
||||||
|
|
||||||
|
|
||||||
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-4)
|
|
||||||
lr_scheduler = StepLR(opt, step_size=3, gamma=0.1)
|
|
||||||
|
|
||||||
criterion = nn.CrossEntropyLoss()
|
|
||||||
|
|
||||||
# Step 3: Create the ART classifier
|
|
||||||
|
|
||||||
classifier = PyTorchClassifier(
|
|
||||||
model=model,
|
|
||||||
clip_values=(0.0, 1.0),
|
|
||||||
loss=criterion,
|
|
||||||
optimizer=opt,
|
|
||||||
input_shape=(3, size, size),
|
|
||||||
nb_classes=10,
|
|
||||||
)
|
|
||||||
|
|
||||||
attack = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=10,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
x_test_clean_pred=np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on clean samples before adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_clean_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# Step 4: Create the trainer object - AdversarialTrainerTRADESPyTorch
|
|
||||||
trainer = AdversarialTrainer(
|
|
||||||
classifier, attack
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build a Keras image augmentation object and wrap it in ART
|
|
||||||
art_datagen = PyTorchDataGenerator(iterator=train_dataloader, size=train_size, batch_size=128)
|
|
||||||
|
|
||||||
# Step 5: fit the trainer
|
|
||||||
trainer.fit_generator(art_datagen, nb_epochs=50)
|
|
||||||
|
|
||||||
|
|
||||||
x_test_pred = np.argmax(classifier.predict(x_test), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on benign test samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
|
|
||||||
attack_test = ProjectedGradientDescent(
|
|
||||||
classifier,
|
|
||||||
norm=np.inf,
|
|
||||||
eps=8.0 / 255.0,
|
|
||||||
eps_step=2.0 / 255.0,
|
|
||||||
max_iter=20,
|
|
||||||
targeted=False,
|
|
||||||
num_random_init=1,
|
|
||||||
batch_size=128,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
x_test_attack = attack_test.generate(x_test, y=y_test)
|
|
||||||
x_test_attack_pred = np.argmax(classifier.predict(x_test_attack), axis=1)
|
|
||||||
print(
|
|
||||||
"Accuracy on original PGD adversarial samples after adversarial training: %.2f%%"
|
|
||||||
% (np.sum(x_test_attack_pred == np.argmax(y_test, axis=1)) / x_test.shape[0] * 100)
|
|
||||||
)
|
|
||||||
torch.save(trainer.classifier.model.state_dict(), 'svhn-pgd.pth')
|
|
||||||
print(
|
|
||||||
"Save the AT model! "
|
|
||||||
)
|
|
||||||
Loading…
Reference in New Issue