def test_adversary_mnist(algorithm, step, p, model_filename):
import os
from advertorch.test_utils import LeNet5
from advertorch_examples.utils import get_mnist_test_loader
from advertorch_examples.utils import TRAINED_MODEL_PATH
# Setup
torch.manual_seed(0)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
device = torch.device("cpu")
model = LeNet5()
model_path = os.path.join(TRAINED_MODEL_PATH, model_filename)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
criterion = nn.CrossEntropyLoss()
constraint = chop.constraints.make_LpBall(alpha=1., p=p)
test_loader = get_mnist_test_loader(batch_size=1, shuffle=True)
# Just get first batch
for data, target in test_loader:
break
data, target = data.to(device), target.to(device)
adv = Adversary(algorithm)
if algorithm == optim.minimize_pgd:
alg_kwargs = {'prox': constraint.prox, 'max_iter': 50, 'step': step}
elif algorithm == optim.minimize_pgd_madry:
alg_kwargs = {
'prox': constraint.prox,
'lmo': constraint.lmo,
'max_iter': 50,
'step': step
}
elif algorithm == optim.minimize_frank_wolfe:
alg_kwargs = {
'lmo': constraint.lmo,
'step': step if step else 'sublinear',
'max_iter': 50
}
# Run and log perturbation
adv_loss, delta = adv.perturb(data, target, model, criterion, **alg_kwargs)
def test_adversary_synthetic_data(algorithm, step_size, p):
# Setup
torch.manual_seed(0)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
data = torch.rand((1, 25, 25))
target = torch.zeros(1).long()
data = data.to(device)
target = target.to(device)
model = LinearModel()
model.to(device)
criterion = nn.CrossEntropyLoss()
constraint = chop.constraints.make_LpBall(alpha=1., p=p)
adv = Adversary(algorithm)
# Get nominal loss
output = model(data)
loss = criterion(output, target)
# Algorithm arguments:
if algorithm == optim.minimize_pgd:
alg_kwargs = {
'prox': constraint.prox,
'max_iter': 50
}
elif algorithm == optim.minimize_pgd_madry:
alg_kwargs = {
'prox': constraint.prox,
'lmo': constraint.lmo,
'max_iter': 50,
'step': 2. * constraint.alpha / 50
}
elif algorithm == optim.minimize_frank_wolfe:
alg_kwargs = {
'lmo': constraint.lmo,
'step': 'sublinear',
'max_iter': 50
}
# Run perturbation
adv_loss, delta = adv.perturb(data, target, model, criterion, **alg_kwargs)
criterion = torch.nn.CrossEntropyLoss()
# TODO: use learning rate schedulers
optimizer = SGD(model.parameters(), lr=.1, momentum=.9, weight_decay=5e-4)
# Define the perturbation constraint set
max_iter_train = 7
max_iter_test = 20
alpha = 8. / 255
constraint = chop.constraints.LinfBall(alpha)
criterion_adv = torch.nn.CrossEntropyLoss(reduction='none')
print(f"Training on L{constraint.p} ball({alpha}).")
adversary = Adversary(chop.optim.minimize_pgd_madry)
for _ in range(n_epochs):
# Train
n_correct = 0
n_correct_adv = 0
model.train()
for k, (data, target) in tqdm(enumerate(trainloader),
total=len(trainloader)):
data = data.to(device)
target = target.to(device)
def image_constraint_prox(delta, step_size=None):
criterion = torch.nn.CrossEntropyLoss(reduction='none')
# Define the perturbation constraint set
max_iter = 20
alpha = 8 / 255.
constraint = chop.constraints.LinfBall(alpha)
print(f"Evaluating model {model_name} on L{constraint.p} ball({alpha}).")
n_correct = 0
n_correct_adv_pgd_madry = 0
n_correct_adv_pgd = 0
n_correct_adv_split = 0
n_correct_adv_fw = 0
adversary_pgd = Adversary(minimize_pgd)
adversary_pgd_madry = Adversary(minimize_pgd_madry)
adversary_split = Adversary(minimize_three_split)
adversary_fw = Adversary(minimize_frank_wolfe)
for k, (data, target) in tqdm(enumerate(loader), total=len(loader)):
data = data.to(device)
target = target.to(device)
def image_constraint_prox(delta, step_size=None):
"""Projects perturbation delta
so that 0. <= data + delta <= 1."""
adv_img = torch.clamp(data + delta, 0, 1)
delta = adv_img - data
return delta