])
trainset = torchvision.datasets.ImageFolder(opt.root+'/sngan_dog_cat', transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = torchvision.datasets.ImageFolder(opt.root+'/sngan_dog_cat_val', transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
else:
raise NotImplementedError('Invalid dataset')
# Model
if opt.model == 'vgg':
from models.vgg import VGG
net = nn.DataParallel(VGG('VGG16', nclass, img_width=img_width).cuda())
elif opt.model == 'aaron':
from models.aaron import Aaron
net = nn.DataParallel(Aaron(nclass).cuda())
else:
raise NotImplementedError('Invalid model')
if opt.resume:
logging.info(f'==> Resume from {opt.model_out}')
net.load_state_dict(torch.load(opt.model_out))
cudnn.benchmark = True
# Loss function
criterion = nn.CrossEntropyLoss()
# Training
def train(epoch):
logging.info('Epoch: %d' % epoch)
img_width=img_width),
device_ids=range(1))
elif opt.model == 'tiny':
if opt.defense in ('plain', 'adv'):
from models.tiny import Tiny
net = nn.DataParallel(Tiny(nclass), device_ids=range(1))
elif opt.defense in ('vi', 'adv_vi'):
from models.tiny_vi import Tiny
net = nn.DataParallel(Tiny(1.0, 1.0, 1.0, nclass), device_ids=range(1))
elif opt.defense in ('rse'):
from models.tiny_rse import Tiny
net = nn.DataParallel(Tiny(nclass, 0.2, 0.1), device_ids=range(1))
elif opt.model == 'aaron':
if opt.defense in ('plain', 'adv'):
from models.aaron import Aaron
net = nn.DataParallel(Aaron(nclass), device_ids=range(1))
elif opt.defense in ('vi', 'adv_vi'):
from models.aaron_vi import Aaron
net = nn.DataParallel(Aaron(1.0, 1.0, 1.0, nclass),
device_ids=range(1))
elif opt.defense in ('rse'):
from models.aaron_rse import Aaron
net = nn.DataParallel(Aaron(nclass, 0.2, 0.1), device_ids=range(1))
else:
raise ValueError('invalid opt.model')
net.load_state_dict(
torch.load(f'./checkpoint/{opt.data}_{opt.model}_{opt.defense}.pth'))
net.cuda()
net.eval() # must set to evaluation mode
loss_f = nn.CrossEntropyLoss()
softmax = nn.Softmax(dim=1)