add cifar100.py

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! "
+)