def get_attack(model, num_classes, attack_name, attack_params):
if attack_name == 'carlini':
attack = attacks.CarliniWagnerL2Attack(model, num_classes,
**attack_params)
elif attack_name == 'elasticnet':
attack = attacks.ElasticNetL1Attack(model, num_classes,
**attack_params)
elif attack_name == 'pgd':
norm = attack_params['norm']
del attack_params['norm']
def get_attack(model, num_classes, attack_name, attack_params):
if attack_name == 'carlini':
attack = attacks.CarliniWagnerL2Attack(model, num_classes,
**attack_params)
elif attack_name == 'elasticnet':
attack = attacks.ElasticNetL1Attack(model, num_classes,
**attack_params)
elif attack_name == 'pgd':
norm = attack_params['norm']
del attack_params['norm']
if norm == 'inf':
attack = attacks.LinfPGDAttack(model, **attack_params)
elif norm == 'l1':
attack = attacks.SparseL1PGDAttack(model, **attack_params)
elif norm == 'l2':
attack = attacks.L2PGDAttack(model, **attack_params)
else:
raise ValueError("Norm not recognized for PGD attack.")
elif attack_name == 'fgsm':
attack = GradientSignAttack(model, **attack_params)
else:
raise ValueError("Attack name not recognized for adv training.")
return attack
def main(init_file, path_model="model_test/blabla",
dataset='ImageNet', num_classes=1000,
epochs=200, batch_size=64,
resume_epoch=0, save_frequency=2,
adversarial_training=None, attack_list=["PGDLinf", "PGDL2"],
eot_samples=1,
noise=None, sigma=0.25):
torch.manual_seed(1234)
job_env = submitit.JobEnvironment()
print(job_env)
torch.cuda.set_device(job_env.local_rank)
torch.distributed.init_process_group(
backend="nccl",
init_method=init_file,
world_size=job_env.num_tasks,
rank=job_env.global_rank,
)
print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
if not os.path.exists(path_model):
os.makedirs(path_model)
# Load inputs
if dataset == "ImageNet":
train_loader = load_data(dataset=dataset,
datadir="/datasets01_101/imagenet_full_size/061417/", # to adapt
batch_size_per_gpu=int(batch_size/job_env.num_tasks),
job_env=job_env, train_mode=True)
else:
train_loader = load_data(dataset=dataset, datadir="datasets",
batch_size_per_gpu=int(batch_size/job_env.num_tasks),
job_env=job_env, train_mode=True)
num_images = len(train_loader.dataset)
# Classifier definition
if dataset == "ImageNet":
# Classifier = models.resnet18(pretrained=False)
Classifier, modelname = getNetwork(net_type='inceptionresnetv2', num_classes=num_classes)
else:
Classifier, modelname = getNetwork(net_type="wide-resnet", depth=28, widen_factor=10,
dropout=0.3, num_classes=num_classes)
Classifier.apply(conv_init)
Classifier = RandModel(Classifier, noise=noise, sigma=sigma)
Classifier.cuda(job_env.local_rank)
cudnn.benchmark = True
Classifier = torch.nn.parallel.DistributedDataParallel(
Classifier, device_ids=[job_env.local_rank], output_device=job_env.local_rank)
Classifier.train()
print("Classifier initialized")
# optimizer and criterion
if adversarial_training == "MixMax":
criterion = torch.nn.CrossEntropyLoss(reduction="none").cuda(job_env.local_rank)
else:
criterion = torch.nn.CrossEntropyLoss().cuda(job_env.local_rank)
optimizer = torch.optim.SGD(
Classifier.parameters(), lr=0.1*batch_size/256, momentum=0.9, weight_decay=5e-4)
if dataset != "ImageNet":
scheduler = get_scheduler(optimizer, policy="multistep", milestones=[
60, 120, 160], gamma=0.2)
else:
scheduler = get_scheduler(optimizer, policy="multistep", milestones=[
30, 60, 90], gamma=0.2)
# resume learning
if resume_epoch > 0:
if os.path.isfile(path_model):
print("=> loading checkpoint '{}'".format(path_model))
checkpoint = torch.load(path_model)
Classifier = checkpoint['net']
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(path_model))
adversaries = dict()
adversaries["CW"] = attacks.CarliniWagnerL2Attack(Classifier, num_classes,
learning_rate=0.01, binary_search_steps=9,
max_iterations=15, abort_early=True,
initial_const=0.001, clip_min=0.0, clip_max=1.)
adversaries["EAD"] = attacks.ElasticNetL1Attack(Classifier, num_classes,
confidence=0,
targeted=False, learning_rate=0.01,
binary_search_steps=9, max_iterations=60,
abort_early=True, initial_const=1e-3,
clip_min=0., clip_max=1., beta=1e-3, decision_rule='EN')
adversaries["PGDL1"] = attacks.SparseL1PGDAttack(Classifier, eps=10., nb_iter=10, eps_iter=2*10./10,
rand_init=False, clip_min=0.0, clip_max=1.0,
sparsity=0.05, eot_samples=eot_samples)
adversaries["PGDLinf"] = attacks.LinfPGDAttack(Classifier, eps=0.031, nb_iter=10, eps_iter=2*0.031/10,
rand_init=True, clip_min=0.0, clip_max=1.0, eot_samples=eot_samples)
adversaries["PGDL2"] = attacks.L2PGDAttack(Classifier, eps=2., nb_iter=10, eps_iter=2*2./10,
rand_init=True, clip_min=0.0, clip_max=1.0, eot_samples=eot_samples)
adversaries["FGSM"] = attacks.GradientSignAttack(Classifier, loss_fn=None, eps=0.05, clip_min=0.,
clip_max=1., targeted=False, eot_samples=eot_samples)
# TO add L1 attacks
for epoch in range(epochs):
current_num_input = 0
running_loss = 0.0
running_acc = 0
start_time_epoch = time.time()
for i, data in enumerate(train_loader, 0):
inputs, labels = data
inputs, labels = inputs.cuda(job_env.local_rank), labels.cuda(job_env.local_rank)
if adversarial_training is None:
outputs = Classifier(inputs)
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if adversarial_training == "Single":
inputs_adv = adversaries[attack_list[0]].perturb(
inputs, labels)
outputs = Classifier(inputs_adv)
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
elif adversarial_training == "MixMean":
loss = 0
for att in attack_list:
inputs_adv = adversaries[att].perturb(inputs, labels)
outputs = Classifier(inputs_adv)
loss += criterion(outputs, labels)
loss /= len(attack_list)
optimizer.zero_grad()
loss.backward()
optimizer.step()
elif adversarial_training == "MixRand":
att = random.choice(attack_list)
inputs_adv = adversaries[att].perturb(inputs, labels)
outputs = Classifier(inputs_adv)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
elif adversarial_training == "MixMax":
loss = torch.zeros_like(labels).float()
for att in attack_list:
inputs_adv = adversaries[att].perturb(inputs, labels)
outputs = Classifier(inputs_adv)
l = criterion(outputs, labels).float()
loss = torch.max(loss, l)
loss = loss.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
outputs = Classifier(inputs)
_, predicted = torch.max(outputs.data, 1)
running_loss += loss.item()
running_acc += predicted.eq(labels.data).cpu().sum().numpy()
curr_batch_size = inputs.size(0)
if i % 5 == 4
print("Epoch :[", epoch+1, "/", epochs,
"] [", i*batch_size, "/", num_images,
"] Running loss:", running_loss/5,
", Running accuracy:", running_acc/(5*curr_batch_size), " time:", time.time()-start_time_epoch)
running_loss = 0.0
running_acc = 0
# save model
if ((epoch + 1) % save_frequency == 0) and (job_env.global_rank == 0):
state = {
'epoch': epoch + 1,
'model_state_dict': Classifier.state_dict(),
}
torch.save(state, os.path.join(
path_model, "epoch_"+str(epoch+1)+'.t7'))
scheduler.step()
if __name__ == "__main__":
main()
def main(
init_file,
path_model="test.t7",
result_file='blabla.txt',
dataset="CIFAR10",
num_classes=10,
batch_size=256,
attack="PGDL2",
eot_samples=1, # 80
noise=None,
batch_prediction=1,
sigma=0.25,
save_image=False):
torch.manual_seed(1234)
job_env = submitit.JobEnvironment()
print(job_env)
torch.cuda.set_device(job_env.local_rank)
torch.distributed.init_process_group(
backend="nccl",
init_method=init_file,
world_size=job_env.num_tasks,
rank=job_env.global_rank,
)
if noise is None:
batch_prediction = None
sigma = None
# Load inputs
test_loader = load_data(dataset=dataset,
datadir="datasets",
batch_size_per_gpu=int(batch_size /
job_env.num_tasks),
job_env=job_env,
train_mode=False)
num_images = len(test_loader.dataset)
# Classifier definition
# torch.nn.Module.dump_patches = True
# model_load = torch.load(path_model)
# Classifier = model_load["net"]
ckpt = torch.load(path_model)
epoch = ckpt["epoch"]
model, _ = getNetwork(net_type="wide-resnet",
depth=28,
widen_factor=10,
dropout=0.3,
num_classes=num_classes)
model.load_state_dict(ckpt["model_state_dict"])
Classifier = RandModel(model, noise=noise, sigma=sigma)
Classifier.cuda(job_env.local_rank)
cudnn.benchmark = True
Classifier = torch.nn.parallel.DistributedDataParallel(
Classifier,
device_ids=[job_env.local_rank],
output_device=job_env.local_rank)
print("Classifier intialized")
for i in range(torch.cuda.device_count()):
print(torch.cuda.get_device_name(i))
# print(Classifier)
Classifier.eval()
adversaries = dict()
adversaries["CW"] = attacks.CarliniWagnerL2Attack(Classifier,
num_classes,
learning_rate=0.01,
binary_search_steps=9,
max_iterations=60,
abort_early=True,
initial_const=0.001,
clip_min=0.0,
clip_max=1.)
adversaries["EAD"] = attacks.ElasticNetL1Attack(Classifier,
num_classes,
confidence=0,
targeted=False,
learning_rate=0.01,
binary_search_steps=9,
max_iterations=60,
abort_early=True,
initial_const=1e-3,
clip_min=0.,
clip_max=1.,
beta=1e-3,
decision_rule='EN')
adversaries["PGDL1"] = attacks.SparseL1PGDAttack(Classifier,
eps=10.,
nb_iter=40,
eps_iter=2 * 10. / 40,
rand_init=False,
clip_min=0.0,
clip_max=1.0,
sparsity=0.05,
eot_samples=eot_samples)
adversaries["PGDLinf"] = attacks.LinfPGDAttack(Classifier,
eps=0.031,
nb_iter=40,
eps_iter=2 * 0.031 / 40,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
eot_samples=eot_samples)
adversaries["PGDL2"] = attacks.L2PGDAttack(Classifier,
eps=2.,
nb_iter=40,
eps_iter=2 * 2. / 40,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
eot_samples=eot_samples)
adversaries["FGSM"] = attacks.GradientSignAttack(Classifier,
loss_fn=None,
eps=0.05,
clip_min=0.,
clip_max=1.,
targeted=False,
eot_samples=eot_samples)
current_num_input = 0
running_acc = 0
if attack is not None:
norms_l1 = []
norms_l2 = []
norms_linf = []
for i, data in enumerate(test_loader, 0):
if i > 0 and save_image:
break
# get the inputs
inputs, labels = data
inputs, labels = inputs.cuda(job_env.local_rank), labels.cuda(
job_env.local_rank)
if (i == 0) and save_image and (job_env.global_rank == 0):
torchvision.utils.save_image(inputs,
"images_nat.jpg",
nrow=8,
padding=2,
normalize=False,
range=None,
scale_each=False,
pad_value=0)
if attack is not None:
inputs_adv = adversaries[attack].perturb(inputs, labels)
norms_l1_batch = get_lp_norm(inputs_adv - inputs, p=1)
norms_l2_batch = get_lp_norm(inputs_adv - inputs, p=2)
norms_linf_batch = get_lp_norm(inputs_adv - inputs, p=np.inf)
norms_l1.append(norms_l1_batch)
norms_l2.append(norms_l2_batch)
norms_linf.append(norms_linf_batch)
inputs = inputs_adv
if (i == 0) and save_image and (job_env.global_rank == 0):
torchvision.utils.save_image(inputs,
"images_adv.jpg",
nrow=8,
padding=2,
normalize=False,
range=None,
scale_each=False,
pad_value=0)
with torch.no_grad():
if noise is None:
outputs = Classifier(inputs)
_, predicted = torch.max(outputs.data, 1)
else:
outputs = torch.FloatTensor(labels.shape[0],
num_classes).cuda()
outputs.zero_()
for _ in range(batch_prediction):
outputs += Classifier(inputs)
_, predicted = torch.max(outputs.data, 1)
# print statistics
running_acc += predicted.eq(labels.data).cpu().sum().numpy()
curr_batch_size = inputs.size(0)
current_num_input += curr_batch_size
print("[", (i + 1) * batch_size, "/", num_images, "] running_acc=",
running_acc / current_num_input)
running_acc = torch.Tensor([running_acc]).cuda(job_env.local_rank)
torch.distributed.all_reduce(running_acc,
op=torch.distributed.ReduceOp.SUM)
accuracy = (running_acc / num_images).cpu().sum().numpy()
print(accuracy)
if attack is not None:
norms_l1 = torch.cat(norms_l1).view(-1)
norms_l2 = torch.cat(norms_l2).view(-1)
norms_linf = torch.cat(norms_linf).view(-1)
norms_l1_gathered = all_gather(norms_l1)
norms_l2_gathered = all_gather(norms_l2)
norms_linf_gathered = all_gather(norms_linf)
norms_l1_gathered = torch.cat(norms_l1_gathered).view(
-1).detach().cpu().numpy()
norms_l2_gathered = torch.cat(norms_l2_gathered).view(
-1).detach().cpu().numpy()
norms_linf_gathered = torch.cat(norms_linf_gathered).view(
-1).detach().cpu().numpy()
if job_env.global_rank == 0:
if attack is not None:
np.save(result_file + "_" + attack + "_l1norm", norms_l1_gathered)
np.save(result_file + "_" + attack + "_l2norm", norms_l2_gathered)
np.save(result_file + "_" + attack + "_linfnorm",
norms_linf_gathered)
with open(result_file + ".txt", 'a') as f:
f.write('{} {} {} {} {} {} {}\n'.format(epoch, dataset, noise,
batch_prediction, attack,
eot_samples, accuracy))
torch.distributed.barrier()
torch.distributed.destroy_process_group()
print(job_env.local_rank, job_env.global_rank)
return job_env.local_rank, job_env.global_rank