添加 mlp_mix.py

auto_cifar10.py

@@ -0,0 +1,132 @@
+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')
+

cifar10.py

@@ -0,0 +1,192 @@
+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! "
+)

cifar100.py

@@ -0,0 +1,182 @@
+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! "
+)

mlp_mix.py

@@ -0,0 +1,92 @@
+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]

svhn_at.py

@@ -0,0 +1,182 @@
+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! "
+)