def poisoned_pixel_CDP(norm_bound, noise_scale, nb_attackers, seed=1):
start_time = time.time()
# define paths
path_project = os.path.abspath('..')
logger = SummaryWriter('../logs')
args = args_parser()
exp_details(args)
# set seed
torch.manual_seed(seed)
np.random.seed(seed)
# device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load dataset and user groups
train_dataset, test_dataset, user_groups = get_dataset(args)
# BUILD MODEL
if args.model == 'cnn':
# Convolutional neural netork
if args.dataset == 'mnist':
global_model = CNNMnist(args=args)
elif args.dataset == 'fmnist':
global_model = CNNFashion_Mnist(args=args)
elif args.dataset == 'cifar':
global_model = CNNCifar(args=args)
elif args.model == 'mlp':
# Multi-layer preceptron
img_size = train_dataset[0][0].shape
len_in = 1
for x in img_size:
len_in *= x
global_model = MLP(dim_in=len_in,
dim_hidden=64,
dim_out=args.num_classes)
else:
exit('Error: unrecognized model')
# Set the model to train and send it to device.
global_model.to(device)
global_model.train()
print(global_model)
# copy weights
global_weights = global_model.state_dict()
# load poisoned model
backdoor_model = copy.deepcopy(global_model)
backdoor_model.load_state_dict(torch.load('../save/poison_model.pth'))
# testing accuracy for global model
testing_accuracy = [0.1]
backdoor_accuracy = [0.1]
for epoch in tqdm(range(args.epochs)):
local_del_w, local_norms = [], []
print(f'\n | Global Training Round : {epoch + 1} |\n')
global_model.train()
m = max(int(args.frac * args.num_users), 1)
idxs_users = np.random.choice(range(args.num_users), m, replace=False)
# Adversary updates
print("Evil")
for idx in idxs_users[0:nb_attackers]:
# backdoor model
w = copy.deepcopy(backdoor_model)
# compute change in parameters and norm
zeta = 0
for del_w, w_old in zip(w.parameters(), global_model.parameters()):
del_w.data = del_w.data - copy.deepcopy(w_old.data)
zeta += del_w.norm(2).item()**2
zeta = zeta**(1. / 2)
del_w = w.state_dict()
print("EVIL")
print(zeta)
# add to global round
local_del_w.append(copy.deepcopy(del_w))
local_norms.append(copy.deepcopy(zeta))
# Non-adversarial updates
for idx in idxs_users[nb_attackers:]:
local_model = LocalUpdate(args=args,
dataset=train_dataset,
idxs=user_groups[idx],
logger=logger)
del_w, zeta = local_model.update_weights(
model=copy.deepcopy(global_model), change=1)
local_del_w.append(copy.deepcopy(del_w))
local_norms.append(copy.deepcopy(zeta))
print("good")
#print(zeta)
# norm bound (e.g. median of norms)
clip_factor = norm_bound #min(norm_bound, np.median(local_norms))
print(clip_factor)
# clip updates
for i in range(len(idxs_users)):
for param in local_del_w[i].values():
print(max(1, local_norms[i] / clip_factor))
param /= max(1, local_norms[i] / clip_factor)
# average local model updates
average_del_w = average_weights(local_del_w)
# Update model and add noise
# w_{t+1} = w_{t} + avg(del_w1 + del_w2 + ... + del_wc) + Noise
for param, param_del_w in zip(global_weights.values(),
average_del_w.values()):
param += param_del_w
param += torch.randn(
param.size()) * noise_scale * norm_bound / len(idxs_users)
global_model.load_state_dict(global_weights)
# test accuracy
test_acc, test_loss, backdoor = test_backdoor_pixel(
args, global_model, test_dataset)
testing_accuracy.append(test_acc)
backdoor_accuracy.append(backdoor)
print("Testing & Backdoor accuracies")
print(testing_accuracy)
print(backdoor_accuracy)
# save test accuracy
np.savetxt(
'../save/PixelAttack/TestAcc/iid_GDP_{}_{}_clip{}_scale{}_attackers{}_seed{}.txt'
.format(args.dataset, args.model, norm_bound, noise_scale,
nb_attackers, s), testing_accuracy)
np.savetxt(
'../save/PixelAttack/BackdoorAcc/iid_GDP_{}_{}_clip{}_scale{}_attackers{}_seed{}.txt'
.format(args.dataset, args.model, norm_bound, noise_scale,
nb_attackers, s), backdoor_accuracy)
len_in *= x
global_model = MLP(dim_in=len_in,
dim_hidden=64,
dim_out=args.num_classes)
else:
exit('Error: unrecognized model')
# Set the model to train and send it to device.
global_model.to(device)
global_model.train()
print(global_model)
# Training
# Set optimizer and criterion
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(global_model.parameters(),
lr=args.lr,
momentum=0.5)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(global_model.parameters(),
lr=args.lr,
weight_decay=1e-4)
trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
criterion = torch.nn.NLLLoss().to(device)
epoch_loss = []
for epoch in tqdm(range(args.epochs)):
batch_loss = []
for batch_idx, (images, labels) in enumerate(trainloader):
def main():
model_path = 'results/%s/%s/%s/seed_%d' % (args.dataset, args.method, args.net_type, args.seed)
if not os.path.isdir(model_path):
mkdir_p(model_path)
# load datasets
train_dataset, test_dataset, _ = get_dataset(args)
# BUILD MODEL
if args.model == 'cnn':
# Convolutional neural netork
if args.dataset == 'mnist':
global_model = CNNMnist(args=args)
elif args.dataset == 'fmnist':
global_model = CNNFashion_Mnist(args=args)
elif args.dataset == 'cifar10':
global_model = CNNCifar(args=args)
elif args.dataset == 'cub200':
if args.net_type == 'resnet':
#global_model = models.resnet50(pretrained=True)
global_model = models.resnet18(pretrained=True)
global_model.fc = torch.nn.Linear(global_model.fc.in_features, cf.num_classes[args.dataset])
elif args.model == 'mlp':
# Multi-layer preceptron
#img_size is torch.Size([1, 28, 28])
img_size = train_dataset[0][0].shape
len_in = 1
#toclarify: why do we have to call the MLP code 3 times?
#TODO: try to move global_model out of the bracket
for x in img_size:
len_in *= x
global_model = MLP(dim_in=len_in, dim_hidden=64,
dim_out=args.num_classes)
else:
exit('Error: unrecognized model')
# Set the model to train and send it to device.
global_model.to(device)
global_model.train()
print(global_model)
wandb.watch(global_model)
# Training
# Set optimizer and criterion
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(global_model.parameters(), lr=args.lr,
momentum=cf.momentum[args.dataset], weight_decay=5e-4)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(global_model.parameters(), lr=args.lr,
weight_decay=1e-4)
#batch_size =
trainloader = DataLoader(train_dataset, batch_size=int(args.local_bs * (args.num_users * args.frac)), shuffle=True, num_workers=args.workers,
pin_memory=use_cuda, drop_last=True)
criterion = torch.nn.CrossEntropyLoss().to(device)
epoch_loss = []
test_acc_lst = []
best_acc = 0
args.lr = cf.lr[args.dataset]
for epoch in tqdm(range(args.epochs)):
global_model.train()
batch_loss = []
# adjest learning rate per global round
if epoch != 0:
adjust_learning_rate([optimizer], args, epoch)
for batch_idx, (images, labels) in enumerate(trainloader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = global_model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
#if batch_idx % 50 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLr: {:.4f}'.format(
epoch+1, (batch_idx+1) * len(images), len(trainloader.dataset),
100. * (batch_idx+1) / len(trainloader), loss.item(), args.lr))
batch_loss.append(loss.item())
wandb.log({'Train Loss': loss.item()})
loss_avg = sum(batch_loss)/len(batch_loss)
print('\nTrain loss: \n', loss_avg)
epoch_loss.append(loss_avg)
test_acc, test_loss = test_inference(args, global_model, test_dataset)
test_acc_lst.append(test_acc)
#save model
if test_acc > best_acc:
best_acc = test_acc
save_checkpoint({
'epoch': epoch + 1,
'state_dict': global_model.state_dict(),
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
}, dir=model_path, filename='checkpoint.pth.tar')
print('\nTrain Epoch: {}, Test acc: {:.2f}%, Best Test acc: {:.2f}%'.format(epoch + 1, test_acc, best_acc))
# log training loss, test accuracy at wandb
wandb.log({'Test Acc': test_acc,
'Best Acc': best_acc})
# if model achieves target test acc, stop training
if best_acc >= args.target_acc:
print('Total Global round: ', epoch+1)
break
if not os.path.isdir(os.path.join(model_path, 'save')):
mkdir_p(os.path.join(model_path, 'save'))
# Plot loss
plt.figure()
plt.plot(range(len(epoch_loss)), epoch_loss)
plt.xlabel('epochs')
plt.ylabel('Train loss')
plt.savefig(os.path.join(model_path, 'save/nn_{}_{}_{}_loss.png'.format(args.dataset, args.model,
args.epochs)))
# Plot test acc per epoch
plt.figure()
plt.plot(range(len(test_acc_lst)), test_acc_lst)
plt.xlabel('epochs')
plt.ylabel('Test accuracy')
plt.savefig(os.path.join(model_path, 'save/nn_{}_{}_{}_acc.png'.format(args.dataset, args.model,
args.epochs)))
# testing
#test_acc, test_loss = test_inference(args, global_model, test_dataset)
print('Test on', len(test_dataset), 'samples')
print("Best Test Accuracy: {:.2f}%".format(best_acc))
def main_test(args):
start_time = time.time()
now = datetime.datetime.now().strftime('%Y-%m-%d-%H%M%S')
# define paths
logger = SummaryWriter('../logs')
# easydict 사용하는 경우 주석처리
# args = args_parser()
# checkpoint 생성위치
args.save_path = os.path.join(args.save_path, args.exp_folder)
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
save_path_tmp = os.path.join(args.save_path, 'tmp_{}'.format(now))
if not os.path.exists(save_path_tmp):
os.makedirs(save_path_tmp)
SAVE_PATH = os.path.join(args.save_path, '{}_{}_T[{}]_C[{}]_iid[{}]_E[{}]_B[{}]'.
format(args.dataset, args.model, args.epochs, args.frac, args.iid,
args.local_ep, args.local_bs))
# 시드 고정
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
# torch.cuda.set_device(0)
device = torch.device("cuda:{}".format(args.gpu) if torch.cuda.is_available() else "cpu")
cpu_device = torch.device('cpu')
# log 파일 생성
log_path = os.path.join('../logs', args.exp_folder)
if not os.path.exists(log_path):
os.makedirs(log_path)
loggertxt = get_logger(
os.path.join(log_path, '{}_{}_{}_{}.log'.format(args.model, args.optimizer, args.norm, now)))
logging.info(args)
# csv
csv_save = '../csv/' + now
csv_path = os.path.join(csv_save, 'accuracy.csv')
csv_logger_keys = ['train_loss', 'accuracy']
csvlogger = CSVLogger(csv_path, csv_logger_keys)
# load dataset and user groups
train_dataset, test_dataset, client_loader_dict = get_dataset(args)
# cifar-100의 경우 자동 설정
if args.dataset == 'cifar100':
args.num_classes = 100
# BUILD MODEL
if args.model == 'cnn':
# Convolutional neural network
if args.dataset == 'mnist':
global_model = CNNMnist(args=args)
elif args.dataset == 'fmnist':
global_model = CNNFashion_Mnist(args=args)
elif args.dataset == 'cifar':
global_model = CNNCifar(args=args)
elif args.dataset == 'cifar100':
global_model = CNNCifar(args=args)
elif args.model == 'mlp':
# Multi-layer preceptron
img_size = train_dataset[0][0].shape
len_in = 1
for x in img_size:
len_in *= x
global_model = MLP(dim_in=len_in, dim_hidden=64,
dim_out=args.num_classes)
elif args.model == 'cnn_vc':
global_model = CNNCifar_fedVC(args=args)
elif args.model == 'cnn_vcbn':
global_model = CNNCifar_VCBN(args=args)
elif args.model == 'cnn_vcgn':
global_model = CNNCifar_VCGN(args=args)
elif args.model == 'resnet18_ws':
global_model = resnet18(num_classes=args.num_classes, weight_stand=1)
elif args.model == 'resnet18':
global_model = resnet18(num_classes=args.num_classes, weight_stand=0)
elif args.model == 'resnet32':
global_model = ResNet32_test(num_classes=args.num_classes)
elif args.model == 'resnet18_mabn':
global_model = resnet18_mabn(num_classes=args.num_classes)
elif args.model == 'vgg':
global_model = vgg11()
elif args.model == 'cnn_ws':
global_model = CNNCifar_WS(args=args)
else:
exit('Error: unrecognized model')
# Set the model to train and send it to device.
loggertxt.info(global_model)
# fedBN처럼 gn no communication 용
client_models = [copy.deepcopy(global_model) for idx in range(args.num_users)]
# copy weights
global_weights = global_model.state_dict()
global_model.to(device)
global_model.train()
# Training
train_loss, train_accuracy = [], []
val_acc_list, net_list = [], []
# how does help BN 확인용
client_loss = [[] for i in range(args.num_users)]
client_conv_grad = [[] for i in range(args.num_users)]
client_fc_grad = [[] for i in range(args.num_users)]
client_total_grad_norm = [[] for i in range(args.num_users)]
# 전체 loss 추적용 -how does help BN
# 재시작
if args.resume:
checkpoint = torch.load(SAVE_PATH)
global_model.load_state_dict(checkpoint['global_model'])
if args.hold_normalize:
for client_idx in range(args.num_users):
client_models[client_idx].load_state_dict(checkpoint['model_{}'.format(client_idx)])
else:
for client_idx in range(args.num_users):
client_models[client_idx].load_state_dict(checkpoint['global_model'])
resume_iter = int(checkpoint['a_iter']) + 1
print('Resume trainig form epoch {}'.format(resume_iter))
else:
resume_iter = 0
# learning rate scheduler
#scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer, gamma=0.1,step_size=500)
# start training
for epoch in tqdm(range(args.epochs)):
local_weights, local_losses = [], []
if args.verbose:
print(f'\n | Global Training Round : {epoch + 1} |\n')
global_model.train()
m = max(int(args.frac * args.num_users), 1)
idxs_users = np.random.choice(range(args.num_users), m, replace=False)
for idx in idxs_users:
"""
for key in global_model.state_dict().keys():
if args.hold_normalize:
if 'bn' not in key:
client_models[idx].state_dict()[key].data.copy_(global_model.state_dict()[key])
else:
client_models[idx].state_dict()[key].data.copy_(global_model.state_dict()[key])
"""
torch.cuda.empty_cache()
local_model = LocalUpdate(args=args, logger=logger, train_loader=client_loader_dict[idx], device=device)
w, loss, batch_loss, conv_grad, fc_grad, total_gard_norm = local_model.update_weights(
model=copy.deepcopy(global_model), global_round=epoch, idx_user=idx)
local_weights.append(copy.deepcopy(w))
# client의 1 epoch에서의 평균 loss값 ex)0.35(즉, batch loss들의 평균)
local_losses.append(copy.deepcopy(loss))
# 전체 round scheduler
# scheduler.step()
# loss graph용 -> client당 loss값 진행 저장 -> 모두 client별로 저장.
client_loss[idx].append(batch_loss)
client_conv_grad[idx].append(conv_grad)
client_fc_grad[idx].append(fc_grad)
client_total_grad_norm[idx].append(total_gard_norm)
# print(total_gard_norm)
# gn, bn 복사
# client_models[idx].load_state_dict(w)
del local_model
del w
# update global weights
global_weights = average_weights(local_weights, client_loader_dict, idxs_users)
# update global weights
# opt = OptRepo.name2cls('adam')(global_model.parameters(), lr=0.01, betas=(0.9, 0.99), eps=1e-3)
opt = OptRepo.name2cls('sgd')(global_model.parameters(), lr=10, momentum=0.9)
opt.zero_grad()
opt_state = opt.state_dict()
global_weights = aggregation(global_weights, global_model)
global_model.load_state_dict(global_weights)
opt = OptRepo.name2cls('sgd')(global_model.parameters(), lr=10, momentum=0.9)
# opt = OptRepo.name2cls('adam')(global_model.parameters(), lr=0.01, betas=(0.9, 0.99), eps=1e-3)
opt.load_state_dict(opt_state)
opt.step()
loss_avg = sum(local_losses) / len(local_losses)
train_loss.append(loss_avg)
global_model.eval()
# for c in range(args.num_users):
# local_model = LocalUpdate(args=args, dataset=train_dataset,
# idxs=user_groups[idx], logger=logger)
# acc, loss = local_model.inference(model=global_model)
# list_acc.append(acc)
# list_loss.append(loss)
# train_accuracy.append(sum(list_acc)/len(list_acc))
train_accuracy = test_inference(args, global_model, test_dataset, device=device)
val_acc_list.append(train_accuracy)
# print global training loss after every 'i' rounds
# if (epoch+1) % print_every == 0:
loggertxt.info(f' \nAvg Training Stats after {epoch + 1} global rounds:')
loggertxt.info(f'Training Loss : {loss_avg}')
loggertxt.info('Train Accuracy: {:.2f}% \n'.format(100 * train_accuracy))
csvlogger.write_row([loss_avg, 100 * train_accuracy])
if (epoch + 1) % 100 == 0:
tmp_save_path = os.path.join(save_path_tmp, 'tmp_{}.pt'.format(epoch+1))
torch.save(global_model.state_dict(),tmp_save_path)
# Test inference after completion of training
test_acc = test_inference(args, global_model, test_dataset, device=device)
print(' Saving checkpoints to {}...'.format(SAVE_PATH))
if args.hold_normalize:
client_dict = {}
for idx, model in enumerate(client_models):
client_dict['model_{}'.format(idx)] = model.state_dict()
torch.save(client_dict, SAVE_PATH)
else:
torch.save({'global_model': global_model.state_dict()}, SAVE_PATH)
loggertxt.info(f' \n Results after {args.epochs} global rounds of training:')
# loggertxt.info("|---- Avg Train Accuracy: {:.2f}%".format(100*train_accuracy[-1]))
loggertxt.info("|---- Test Accuracy: {:.2f}%".format(100 * test_acc))
# frac이 1이 아닐경우 잘 작동하지않음.
# batch_loss_list = np.array(client_loss).sum(axis=0) / args.num_users
# conv_grad_list = np.array(client_conv_grad).sum(axis=0) / args.num_users
# fc_grad_list = np.array(client_fc_grad).sum(axis=0) / args.num_users
# total_grad_list = np.array(client_total_grad_norm).sum(axis=0) /args.num_users
# client의 avg를 구하고 싶었으나 현재는 client 0만 확인
# client마다 batch가 다를 경우 bug 예상
return train_loss, val_acc_list, client_loss[0], client_conv_grad[0], client_fc_grad[0], client_total_grad_norm[0]
def poisoned_NoDefense(nb_attackers, seed=1):
# define paths
path_project = os.path.abspath('..')
logger = SummaryWriter('../logs')
args = args_parser()
exp_details(args)
# set seed
torch.manual_seed(seed)
np.random.seed(seed)
# device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load dataset and user groups
train_dataset, test_dataset, user_groups = get_dataset(args)
# BUILD MODEL
if args.model == 'cnn':
# Convolutional neural netork
if args.dataset == 'mnist':
global_model = CNNMnist(args=args)
elif args.dataset == 'fmnist':
global_model = CNNFashion_Mnist(args=args)
elif args.dataset == 'cifar':
global_model = CNNCifar(args=args)
elif args.model == 'mlp':
# Multi-layer preceptron
img_size = train_dataset[0][0].shape
len_in = 1
for x in img_size:
len_in *= x
global_model = MLP(dim_in=len_in, dim_hidden=64,
dim_out=args.num_classes)
else:
exit('Error: unrecognized model')
# Set the model to train and send it to device.
global_model.to(device)
global_model.train()
print(global_model)
# copy weights
global_weights = global_model.state_dict()
# backdoor model
dummy_model = copy.deepcopy(global_model)
dummy_model.load_state_dict(torch.load('../save/all_5_model.pth'))
dummy_norm = 0
for x in dummy_model.state_dict().values():
dummy_norm += x.norm(2).item() ** 2
dummy_norm = dummy_norm ** (1. / 2)
# testing accuracy for global model
testing_accuracy = [0.1]
for epoch in tqdm(range(args.epochs)):
local_del_w = []
print(f'\n | Global Training Round : {epoch+1} |\n')
global_model.train()
m = max(int(args.frac * args.num_users), 1)
idxs_users = np.random.choice(range(args.num_users), m, replace=False)
# Adversary updates
for idx in idxs_users[0:nb_attackers]:
print("evil")
local_model = LocalUpdate(args=args, dataset=train_dataset, idxs=user_groups[idx], logger=logger)
#del_w, _ = local_model.poisoned_SGA(model=copy.deepcopy(global_model), change=1)
w = copy.deepcopy(dummy_model)
# compute change in parameters and norm
zeta = 0
for del_w, w_old in zip(w.parameters(), global_model.parameters()):
del_w.data -= copy.deepcopy(w_old.data)
del_w.data *= m / nb_attackers
del_w.data += copy.deepcopy(w_old.data)
zeta += del_w.norm(2).item() ** 2
zeta = zeta ** (1. / 2)
del_w = copy.deepcopy(w.state_dict())
local_del_w.append(copy.deepcopy(del_w))
# Non-adversarial updates
for idx in idxs_users[nb_attackers:]:
print("good")
local_model = LocalUpdate(args=args, dataset=train_dataset, idxs=user_groups[idx], logger=logger)
del_w, _ = local_model.update_weights(model=copy.deepcopy(global_model), change=1)
local_del_w.append(copy.deepcopy(del_w))
# average local updates
average_del_w = average_weights(local_del_w)
# Update global model: w_{t+1} = w_{t} + average_del_w
for param, param_del_w in zip(global_weights.values(), average_del_w.values()):
param += param_del_w
global_model.load_state_dict(global_weights)
# test accuracy
test_acc, test_loss = test_inference(args, global_model, test_dataset)
testing_accuracy.append(test_acc)
print("Test accuracy")
print(testing_accuracy)
# save test accuracy
np.savetxt('../save/RandomAttack/NoDefense_iid_{}_{}_attackers{}_seed{}.txt'.
format(args.dataset, args.model, nb_attackers, s), testing_accuracy)