def prep_loaders():
train_set, test_set = None, None
if DATASET == 'MNIST':
train_set = torchvision.datasets.MNIST(root='.',
train=True,
download=False,
transform=transform_mnist)
test_set = torchvision.datasets.MNIST(root='.',
train=False,
download=False,
transform=transform_mnist)
elif DATASET == 'CIFAR':
train_set = torchvision.datasets.CIFAR10(root='.',
train=True,
download=False,
transform=transform_train)
test_set = torchvision.datasets.CIFAR10(root='.',
train=False,
download=False,
transform=transform_test)
elif DATASET == 'SVHN':
train_set = torchvision.datasets.SVHN(root='.',
split='train',
download=False,
transform=transform_train)
test_set = torchvision.datasets.SVHN(root='.',
split='test',
download=False,
transform=transform_test)
elif DATASET == 'FMNIST':
train_set = torchvision.datasets.FashionMNIST(
root='.', train=True, download=False, transform=transform_mnist)
test_set = torchvision.datasets.FashionMNIST(root='.',
train=False,
download=False,
transform=transform_mnist)
elif DATASET == 'GTSRB':
train_set = dataset.GTSRB(root_dir='.',
train=True,
transform=transform_gtsrb)
test_set = dataset.GTSRB(root_dir='.',
train=False,
transform=transform_gtsrb)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=BATCH_SIZE_TRAIN,
shuffle=True,
num_workers=2)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=BATCH_SIZE_TEST,
shuffle=False,
num_workers=2)
return train_loader, test_loader
def load_data():
if DATASET_NAME == 'cifar':
img_transform_cifar = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
dataset = datasets.CIFAR10(root='CIFAR',
train=True,
transform=img_transform_cifar,
target_transform=None,
download=True)
elif DATASET_NAME == 'gtsrb':
img_transform_gtrsb = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.3403, 0.3121, 0.3214),
(0.2724, 0.2608, 0.2669))
])
dataset = gtsrb_dataset.GTSRB(root_dir='./',
train=True,
transform=img_transform_gtrsb)
elif DATASET_NAME == 'mnist':
img_transform_mnist = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(mean=0.5, std=0.5)
])
dataset = datasets.MNIST(root='./',
train=True,
transform=img_transform_mnist,
download=True)
elif DATASET_NAME == 'mnistm':
img_transform_mnist = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(mean=0.5, std=0.5)
])
train_list = './mnist_m/mnist_m_train_labels.txt'
dataset = GetLoader(data_root='./mnist_m/mnist_m_train',
data_list=train_list,
transform=img_transform_mnist)
else:
print('Data not found.')
exit()
return dataset
def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# Dataset preprocessing
title = 'GTSRB'
# Create Datasets
transform_train_poisoned = transforms.Compose([
transforms.Resize((32, 32)),
TriggerAppending(trigger=args.trigger, alpha=args.alpha),
transforms.ToTensor(),
])
transform_train_benign = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
])
transform_test_poisoned = transforms.Compose([
transforms.Resize((32, 32)),
TriggerAppending(trigger=args.trigger, alpha=args.alpha),
transforms.ToTensor(),
])
transform_test_benign = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
])
print('==> Loading the dataset')
# Create Datasets
num_training = len(
dataset.GTSRB(root_dir='./data',
train=True,
transform=transform_train_benign))
num_poisoned = int(num_training * args.poison_rate)
idx = list(np.arange(num_training))
random.shuffle(idx)
poisoned_idx = idx[:num_poisoned]
benign_idx = idx[num_poisoned:]
poisoned_trainset = dataset.GTSRB_subset(
root_dir='./data',
train=True,
transform=transform_train_poisoned,
List=poisoned_idx,
y_target=args.y_target)
benign_trainset = dataset.GTSRB_subset(root_dir='./data',
train=True,
transform=transform_train_benign,
List=benign_idx,
y_target=None)
poisoned_testset = dataset.GTSRB(root_dir='./data',
train=False,
transform=transform_test_poisoned,
y_target=args.y_target)
benign_testset = dataset.GTSRB(root_dir='./data',
train=False,
transform=transform_test_benign,
y_target=None)
poisoned_trainloader = torch.utils.data.DataLoader(
poisoned_trainset,
batch_size=int(args.train_batch * args.poison_rate),
shuffle=True,
num_workers=args.workers)
benign_trainloader = torch.utils.data.DataLoader(
benign_trainset,
batch_size=int(args.train_batch * (1 - args.poison_rate) * 0.9),
shuffle=True,
num_workers=args.workers
) # *0.9 to prevent the iterations of benign data is less than that of poisoned data
poisoned_testloader = torch.utils.data.DataLoader(
poisoned_testset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers)
benign_testloader = torch.utils.data.DataLoader(benign_testset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers)
print(
"Num of training samples %i (Num of poisoned samples %i, Num of benign samples %i), Num of testing samples %i"
% (num_training, num_poisoned, num_training - num_poisoned,
len(benign_testset)))
# Model
model = ResNet18()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print('Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Resume
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Benign Valid Loss',
'Poisoned Valid Loss', 'Train ACC.', 'Benign Valid ACC.',
'Poisoned Valid ACC.'
])
if args.evaluate:
print('\nEvaluation only')
test_loss, test_acc = test(testloader, model, criterion, start_epoch,
use_cuda)
print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
return
# Train and val
for epoch in range(start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
train_loss, train_acc = train(args, model, poisoned_trainloader,
benign_trainloader, criterion, optimizer,
epoch, use_cuda)
test_loss_benign, test_acc_benign = test(benign_testloader, model,
criterion, epoch, use_cuda)
test_loss_poisoned, test_acc_poisoned = test(poisoned_testloader,
model, criterion, epoch,
use_cuda)
# append logger file
logger.append([
state['lr'], train_loss, test_loss_benign, test_loss_poisoned,
train_acc, test_acc_benign, test_acc_poisoned
])
# save model
is_best = test_acc_benign > best_acc
best_acc = max(test_acc_benign, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'acc': test_acc_benign,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
checkpoint=args.checkpoint)
logger.close()
logger.plot()
savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# Dataset preprocessing
title = 'GTSRB'
print('==> Preparing GTSRB dataset')
transform = transforms.Compose(
[transforms.Resize((32, 32)),
transforms.ToTensor()])
# Create Datasets
trainset = dataset.GTSRB(root_dir='./data',
train=True,
transform=transform)
testset = dataset.GTSRB(root_dir='./data',
train=False,
transform=transform)
# Load Datasets
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers)
testloader = torch.utils.data.DataLoader(testset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers)
# Model
model = vgg19_bn()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print('Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Resume
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
if args.evaluate:
print('\nEvaluation only')
test_loss, test_acc = test(testloader, model, criterion, start_epoch,
use_cuda)
print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
return
# Train and val
for epoch in range(start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
train_loss, train_acc = train(args, model, trainloader, criterion,
optimizer, epoch, use_cuda)
test_loss, test_acc = test(testloader, model, criterion, epoch,
use_cuda)
# append logger file
logger.append(
[state['lr'], train_loss, test_loss, train_acc, test_acc])
# save model
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
checkpoint=args.checkpoint)
logger.close()
logger.plot()
savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
transforms.ToTensor(),
transforms.Normalize((0.3403, 0.3121, 0.3214), (0.2724, 0.2608, 0.2669))
])
#Function to display an image
def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# Create Datasets
trainset = dataset.GTSRB(root_dir='german-traffic-sign',
train=True,
transform=transform)
testset = dataset.GTSRB(root_dir='german-traffic-sign',
train=False,
transform=transform)
# Load Datasets
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=128,
shuffle=True,
num_workers=2)
testloader = torch.utils.data.DataLoader(testset,
batch_size=128,
shuffle=False,
num_workers=2)