File size: 1,980 Bytes
fb9bb41 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F #233
import torch.optim as optim
from torchvision import datasets,models,transforms
from PIL import Image
import argparse
from deeprobust.image.attack.fgsm import FGSM
from deeprobust.image.netmodels.CNN import Net
from deeprobust.image.config import attack_params
from deeprobust.image.utils import download_model
def parameter_parser():
parser = argparse.ArgumentParser(description = "Run attack algorithms.")
parser.add_argument("--destination",
default = './trained_models/',
help = "choose destination to load the pretrained models.")
parser.add_argument("--filename",
default = "MNIST_CNN_epoch_20.pt")
return parser.parse_args()
args = parameter_parser() # read argument and creat an argparse object
model = Net()
model.load_state_dict(torch.load(args.destination + args.filename))
model.eval()
print("Finish loading network.")
xx = datasets.MNIST('./', download = False).data[999:1000].to('cuda')
xx = xx.unsqueeze_(1).float()/255
#print(xx.size())
## Set Target
yy = datasets.MNIST('./', download = False).targets[999:1000].to('cuda')
"""
Generate adversarial examples
"""
F1 = FGSM(model, device = "cuda") ### or cuda
AdvExArray = F1.generate(xx, yy, **attack_params['FGSM_MNIST'])
predict0 = model(xx)
predict0= predict0.argmax(dim=1, keepdim=True)
predict1 = model(AdvExArray)
predict1= predict1.argmax(dim=1, keepdim=True)
print("original prediction:")
print(predict0)
print("attack prediction:")
print(predict1)
xx = xx.cpu().detach().numpy()
AdvExArray = AdvExArray.cpu().detach().numpy()
import matplotlib.pyplot as plt
plt.imshow(xx[0,0]*255,cmap='gray',vmin=0,vmax=255)
plt.savefig('./adversary_examples/mnist_advexample_fgsm_ori.png')
plt.imshow(AdvExArray[0,0]*255,cmap='gray',vmin=0,vmax=255)
plt.savefig('./adversary_examples/mnist_advexample_fgsm_adv.png')
|