CGSCORE / examples /image /test_nattack.py
Yaning1001's picture
Add files using upload-large-folder tool
fb9bb41 verified
import numpy as np
import torch
import torch.nn as nn
from torchvision import datasets, models, transforms
from deeprobust.image.attack.Nattack import NATTACK
from deeprobust.image.netmodels.CNN import Net
#initialize model
model = Net()
model.load_state_dict(torch.load("trained_models/mnist_fgsmtraining_0.2.pt", map_location = torch.device('cuda')))
model.eval()
print("----------model_parameters-----------")
for names,parameters in model.named_parameters():
print(names,',', parameters.type())
print("-------------------------------------")
data_loader = torch.utils.data.DataLoader(
datasets.MNIST('deeprobust/image/data', train = True,
download = True,
transform = transforms.Compose([transforms.ToTensor()])),
batch_size = 1,
shuffle = True)
attack = NATTACK(model)
attack.generate(dataloader = data_loader, classnum = 10)