From e61c116f1d7fec7142c7009ef0893ae8660331b5 Mon Sep 17 00:00:00 2001 From: liwenyun Date: Mon, 28 Oct 2024 19:28:53 +0800 Subject: [PATCH] add cifar100.py --- cifar100.py | 182 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 182 insertions(+) create mode 100644 cifar100.py diff --git a/cifar100.py b/cifar100.py new file mode 100644 index 0000000..8b225d4 --- /dev/null +++ b/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! " +) \ No newline at end of file