File size: 802 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 | import torch
from torchvision import datasets, transforms
import numpy as np
from deeprobust.image.defense.trades import TRADES
from deeprobust.image.netmodels.CNN import Net
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('deeprobust/image/defense/data', train = True, download = True,
transform = transforms.Compose([transforms.ToTensor()])),
batch_size = 100,
shuffle = True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('deeprobust/image/defense/data', train = False,
transform = transforms.Compose([transforms.ToTensor()])),
batch_size = 1000,
shuffle = True)
model = Net()
defense = TRADES(model,'cuda')
defense.generate(train_loader, test_loader)
|