print("skipping target: {}".format(target_idx))
no_save = False
break
if args.device == 'cuda':
state_dict = torch.load(poisons_path)
else:
state_dict = torch.load(poisons_path,
map_location=torch.device('cpu'))
poison_tuple_list, base_idx_list = state_dict[
'poison'], state_dict['idx']
print("Poisons loaded")
poisoned_dset = PoisonedDataset(
args.train_data_path,
subset='others',
transform=transform_train,
num_per_label=args.num_per_class,
poison_tuple_list=poison_tuple_list,
poison_indices=base_idx_list,
subset_group=args.subset_group)
print("Poisoned dataset created")
res[ite] = get_stats(
'{}/{}/log.txt'.format(args.eval_poisons_root, target_idx),
state_dict, ite - 1)
res[ite]['victims'] = {}
for victim_name in args.target_net:
print(victim_name)
victim_net = load_pretrained_net(
victim_name,
args.test_chk_name,
model_chk_path=args.model_resume_path,
def train_network_with_poison(net, target_imgs, targets_indices, poison_tuple_list, base_idx_list, chk_path, args,
save_state=False, eval_targets=None, device='cuda'):
# requires implementing a get_penultimate_params_list() method to get the parameter identifier of the net's last
# layer
if args.end2end:
params = net.parameters()
else:
params = net.module.get_penultimate_params_list()
if args.retrain_opt == 'adam':
print("Using Adam for retraining")
optimizer = torch.optim.Adam(params, lr=args.retrain_lr, weight_decay=args.retrain_wd)
else:
print("Using SGD for retraining")
optimizer = torch.optim.SGD(params, lr=args.retrain_lr, momentum=args.retrain_momentum,
weight_decay=args.retrain_wd)
net.eval()
criterion = nn.CrossEntropyLoss().to('cuda')
# Create the poisoned dataset
transform_train = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
poisoned_dset = PoisonedDataset(args.train_data_path, subset='others', transform=transform_train,
num_per_label=args.num_per_class, poison_tuple_list=poison_tuple_list,
poison_indices=base_idx_list, subset_group=args.subset_group)
poisoned_loader = torch.utils.data.DataLoader(poisoned_dset, batch_size=args.retrain_bsize, shuffle=True)
# The test set of clean CIFAR10
testset = torchvision.datasets.CIFAR10(root=args.dset_path, train=False, download=True, transform=transform_test)
test_loader = torch.utils.data.DataLoader(testset, batch_size=500)
if not args.end2end:
# create a dataloader that returns the features
poisoned_loader = torch.utils.data.DataLoader(FeatureSet(poisoned_loader, net, device=args.device),
batch_size=64, shuffle=True)
for epoch in range(args.retrain_epochs):
net.eval()
loss_meter = AverageMeter()
acc_meter = AverageMeter()
time_meter = AverageMeter()
if epoch in args.lr_decay_epoch:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.1
end_time = time.time()
for ite, (input, target) in enumerate(poisoned_loader):
input, target = input.to('cuda'), target.to('cuda')
if args.end2end:
feat = net.module.penultimate(input)
else:
feat = input
output = net.module.linear(feat)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
prec1 = accuracy(output, target)[0]
time_meter.update(time.time() - end_time)
end_time = time.time()
loss_meter.update(loss.item(), input.size(0))
acc_meter.update(prec1.item(), input.size(0))
if epoch % 30 == 0 and (ite == len(poisoned_loader) - 1):
print("{2}, Epoch {0}, Iteration {1}, loss {loss.val:.3f} ({loss.avg:.3f}), "
"acc {acc.val:.3f} ({acc.avg:.3f})".
format(epoch, ite, time.strftime("%Y-%m-%d %H:%M:%S"),
loss=loss_meter, acc=acc_meter))
sys.stdout.flush()
if epoch == args.retrain_epochs - 1:
# print the scores for target and base
print("------------")
print("Stats after retraining for {} epochs".format(epoch))
poison_pred_list = []
for poison_img, _ in poison_tuple_list:
base_scores = net(poison_img[None, :, :, :].to(device))
base_score, base_pred = base_scores.topk(1, 1, True, True)
poison_pred_list.append(base_pred.item())
print("Target label: {}, Poison label: {}, Poisons' Predictions:{}"
.format(args.target_label, args.poison_label, poison_pred_list))
targets_acc_meter = AverageMeter()
for idx, target_img in zip(targets_indices, target_imgs):
# print the scores for target and base
target_pred = net(target_img.to(device))
score, pred = target_pred.topk(1, 1, True, True)
prec = accuracy(target_pred, torch.tensor([args.poison_label]).to(device))[0]
targets_acc_meter.update(prec.item(), target_img.size(0))
print("Target sample: {}, Prediction:{}, Target's Score:{}".format(idx,
pred[0][0].item(), list(
target_pred.detach().view(-1).cpu().numpy())))
print("Current ATTACK acc on target samples: {acc.avg:.3f}".
format(acc=targets_acc_meter))
# Evaluate the results on the clean test set
print("-----------------------------------------------")
print("Now evaluating the network on the clean test set")
val_acc_meter = AverageMeter()
with torch.no_grad():
for ite, (input, target) in enumerate(test_loader):
input, target = input.to(device), target.to(device)
output = net(input)
prec1 = accuracy(output, target)[0]
val_acc_meter.update(prec1.item(), input.size(0))
if ite % 100 == 0 or ite == len(test_loader) - 1:
print("{2} Epoch {0}, Val iteration {1}, "
"acc {acc.val:.3f} ({acc.avg:.3f})".
format(epoch, ite, time.strftime("%Y-%m-%d %H:%M:%S"), acc=val_acc_meter))
print("* Prec: {}".format(val_acc_meter.avg))
print("-----------------------------")
eval_targets_acc_meter = AverageMeter()
if eval_targets is not None:
print("Now evaluating on the extenral target eval set")
preds = []
with torch.no_grad():
for ite, (input, target) in enumerate(eval_targets):
input, target = input.to(device), target.to(device)
output = net(input)
_, pred = output.topk(1, 1, True, True)
preds.append(str(pred[0][0].item()))
prec1 = accuracy(output, torch.tensor([args.poison_label]).to(device))[0]
eval_targets_acc_meter.update(prec1.item(), input.size(0))
print("Predictions for the external targets: {}".format(", ".join(preds)))
print("ATTACK acc: {}".format(eval_targets_acc_meter.avg))
if save_state:
state_dict = {"state_dict": net.state_dict(), "epoch": epoch, "acc": val_acc_meter.avg}
torch.save(state_dict, os.path.join(chk_path, 'last_epoch.pth'))
return targets_acc_meter.avg, eval_targets_acc_meter.avg