class FL_client():
def __init__(self, args):
if args.dataset == 'cifar':
self.net = CNNCifar(args=args).to(args.device)
else:
self.net = CNNMnist(args=args).to(args.device)
self.net.train()
self.loss_func = nn.CrossEntropyLoss()
self.optimizer = torch.optim.SGD(self.net.parameters(), lr=args.lr)
self.args = args
self.w_glob = []
# key exchange
self.x = self.gx = 0
self.keys = defaultdict(int)
def set_data(self, dataset, idxs):
self.data = DataLoader(DatasetSplit(dataset, idxs),
batch_size=self.args.local_bs,
shuffle=True)
def load_state(self, state_dict):
self.net.load_state_dict(state_dict)
def train(self):
epoch_loss = []
for _ in range(self.args.local_ep):
batch_loss = []
for _, (images, labels) in enumerate(self.data):
images, labels = images.to(self.args.device), labels.to(
self.args.device)
pred = self.net(images)
loss = self.loss_func(pred, labels)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
batch_loss.append(loss.item())
epoch_loss.append(sum(batch_loss) / len(batch_loss))
return self.net.state_dict(), sum(epoch_loss) / len(epoch_loss)
#30 MALICIOUS
w_locals_30.append(copy.deepcopy(w30))
loss_locals_30.append(copy.deepcopy(loss30))
# update global weights
w_glob = FedAvg(w_locals)
w_glob_1 = FedAvg(w_locals_1)
w_glob_5 = FedAvg(w_locals_5)
w_glob_10 = FedAvg(w_locals_10)
w_glob_15 = FedAvg(w_locals_15)
w_glob_20 = FedAvg(w_locals_20)
w_glob_25 = FedAvg(w_locals_25)
w_glob_30 = FedAvg(w_locals_30)
# copy weight to net_glob
net_glob.load_state_dict(w_glob)
net_glob1.load_state_dict(w_glob_1)
net_glob5.load_state_dict(w_glob_5)
net_glob10.load_state_dict(w_glob_10)
net_glob15.load_state_dict(w_glob_15)
net_glob20.load_state_dict(w_glob_20)
net_glob25.load_state_dict(w_glob_25)
net_glob30.load_state_dict(w_glob_30)
# print loss
loss_avg = sum(loss_locals) / len(loss_locals)
loss_avg_1 = sum(loss_locals_1) / len(loss_locals_1)
loss_avg_5 = sum(loss_locals_5) / len(loss_locals_5)
loss_avg_10 = sum(loss_locals_10) / len(loss_locals_10)
loss_avg_15 = sum(loss_locals_15) / len(loss_locals_15)
loss_avg_20 = sum(loss_locals_20) / len(loss_locals_20)
w_locals_30.append(copy.deepcopy(w30))
loss_locals_30.append(copy.deepcopy(loss30))
# update global weights
w_glob = FedAvg(w_locals)
w_glob_1 = FedAvg(w_locals_1)
w_glob_5 = FedAvg(w_locals_5)
w_glob_10 = FedAvg(w_locals_10)
w_glob_15 = FedAvg(w_locals_15)
w_glob_20 = FedAvg(w_locals_20)
w_glob_25 = FedAvg(w_locals_25)
w_glob_30 = FedAvg(w_locals_30)
# copy weight to net_glob
net_glob.load_state_dict(w_glob)
net_glob1.load_state_dict(w_glob_1)
net_glob5.load_state_dict(w_glob_5)
net_glob10.load_state_dict(w_glob_10)
net_glob15.load_state_dict(w_glob_15)
net_glob20.load_state_dict(w_glob_20)
net_glob25.load_state_dict(w_glob_25)
net_glob30.load_state_dict(w_glob_30)
# print loss
loss_avg = sum(loss_locals) / len(loss_locals)
loss_avg_1 = sum(loss_locals_1) / len(loss_locals_1)
loss_avg_5 = sum(loss_locals_5) / len(loss_locals_5)
loss_avg_10 = sum(loss_locals_10) / len(loss_locals_10)
loss_avg_15 = sum(loss_locals_15) / len(loss_locals_15)
loss_avg_20 = sum(loss_locals_20) / len(loss_locals_20)
loss_avg_25 = sum(loss_locals_25) / len(loss_locals_25)
#30 MALICIOUS
w_locals_10.append(copy.deepcopy(w10))
loss_locals_10.append(copy.deepcopy(loss10))
# update global weights
w_glob = FedAvg(w_locals)
w_glob1 = FedAvg(w_locals_1)
w_glob5 = FedAvg(w_locals_5)
w_glob7 = FedAvg(w_locals_7)
w_glob10 = FedAvg(w_locals_10)
# copy weight to net_glob
net_glob.load_state_dict(w_glob)
net_glob1.load_state_dict(w_glob1)
net_glob5.load_state_dict(w_glob5)
net_glob7.load_state_dict(w_glob7)
net_glob10.load_state_dict(w_glob10)
# print loss
loss_avg = sum(loss_locals) / len(loss_locals)
loss_avg_1 = sum(loss_locals_1) / len(loss_locals_1)
loss_avg_5 = sum(loss_locals_5) / len(loss_locals_5)
loss_avg_7 = sum(loss_locals_7) / len(loss_locals_7)
loss_avg_10 = sum(loss_locals_10) / len(loss_locals_10)
print('NO ATTACK ---> Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg))
print('C1 ATTACK ---> Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg_1))
print('C5 ATTACK ---> Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg_5))
def main_worker(gpu, ngpus_per_node, args):
print("gpu:", gpu)
args.gpu = gpu
if args.rank == 0: #(第一台服务器只有三台GPU，需要特殊处理)
newrank = args.rank * ngpus_per_node + gpu
else:
newrank = args.rank * ngpus_per_node + gpu - 1
#初始化,使用tcp方式进行通信
print("begin init")
dist.init_process_group(init_method=args.init_method,
backend="nccl",
world_size=args.world_size,
rank=newrank)
print("end init")
#建立通信group,rank=0作为server，用broadcast模拟send和rec，需要server和每个client建立group
group = []
for i in range(1, args.world_size):
group.append(dist.new_group([0, i]))
allgroup = dist.new_group([i for i in range(args.world_size)])
if newrank == 0:
""" server"""
print("使用{}号服务器的第{}块GPU作为server".format(args.rank, gpu))
#在模型训练期间，server只负责整合参数并分发，不参与任何计算
#设置cpu
args.device = torch.device(
'cuda:{}'.format(args.gpu)
if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
net = CNNMnist().to(args.device)
w_avg = copy.deepcopy(net.state_dict())
for j in range(args.epochs):
if j == args.epochs - 1:
for i in w_avg.keys():
temp = w_avg[i].to(args.device)
w_avg[i] = average_gradients(temp, group, allgroup)
else:
for i in w_avg.keys():
temp = w_avg[i].to(args.device)
average_gradients(temp, group, allgroup)
torch.save(w_avg, 'w_wag')
net.load_state_dict(w_avg)
#加载测试数据
trans_mnist = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])
dataset_test = datasets.MNIST('data/',
train=False,
download=True,
transform=trans_mnist)
test_set = torch.utils.data.DataLoader(dataset_test,
batch_size=args.bs)
test_accuracy, test_loss = test(net, test_set, args)
print("Testing accuracy: {:.2f}".format(test_accuracy))
print("Testing loss: {:.2f}".format(test_loss))
else:
"""clents"""
print("使用{}号服务器的第{}块GPU作为第{}个client".format(args.rank, gpu, newrank))
#设置gpu
args.device = torch.device(
'cuda:{}'.format(args.gpu)
if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
print("begin train...")
net = CNNMnist().to(args.device)
print(net)
data = torch.load("data/distributed/data_of_client{}".format(newrank))
bsz = 64
train_set = torch.utils.data.DataLoader(data, batch_size=bsz)
optimizer = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.5)
num_batches = ceil(len(train_set.dataset) / float(bsz))
start = time.time()
for epoch in range(args.epochs):
for iter in range(3):
epoch_loss = 0.0
for data, target in train_set:
data, target = data.to(args.device), target.to(args.device)
data, target = Variable(data), Variable(target)
optimizer.zero_grad()
output = net(data)
loss = F.nll_loss(output, target)
epoch_loss += loss.item()
loss.backward()
optimizer.step()
if iter == 3 - 1:
print('Rank ', dist.get_rank(), ', epoch ', epoch, ': ',
epoch_loss / num_batches)
"""federated learning"""
w_avg = copy.deepcopy(net.state_dict())
for k in w_avg.keys():
print("k:", k)
temp = average_gradients(w_avg[k].to(args.device), group,
allgroup)
w_avg[k] = temp
net.load_state_dict(w_avg)
end = time.time()
print(" training time:{}".format((end - start)))
train_accuracy, train_loss = test(net, train_set, args)
print("Training accuracy: {:.2f}".format(train_accuracy))
print("Training loss: {:.2f}".format(train_loss))
def main_worker(gpu, ngpus_per_node, args):
print("gpu:", gpu)
args.gpu = gpu
if args.rank == 0: #(第一台服务器只有三台GPU，需要特殊处理)
newrank = args.rank * ngpus_per_node + gpu
else:
newrank = args.rank * ngpus_per_node + gpu - 1
#初始化,使用tcp方式进行通信
dist.init_process_group(init_method=args.init_method,
backend="nccl",
world_size=args.world_size,
rank=newrank)
#建立通信group,rank=0作为server，用broadcast模拟send和rec，需要server和每个client建立group
group = []
for i in range(1, args.world_size):
group.append(dist.new_group([0, i]))
allgroup = dist.new_group([i for i in range(args.world_size)])
if newrank == 0:
""" server"""
print("{}号服务器的第{}块GPU作为server".format(args.rank, gpu))
#在模型训练期间，server只负责整合参数并分发，不参与任何计算
#设置cpu
args.device = torch.device(
'cuda:{}'.format(args.gpu)
if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
#加载测试数据
trans_mnist = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])
dataset_test = datasets.MNIST('data/',
train=False,
download=True,
transform=trans_mnist)
test_set = torch.utils.data.DataLoader(dataset_test,
batch_size=args.bs)
"""calculate influence function"""
model = CNNMnist().to(args.device)
model.load_state_dict(torch.load('w_wag'))
test_id = 0 #选择的test数据id
data, target = test_set.dataset[test_id]
data = test_set.collate_fn([data])
target = test_set.collate_fn([target])
print("begin grad")
grad_test = grad_z(data, target, model, gpu, create_graph=False) #v初始值
print("end grad")
v = grad_test
s_test = []
"""server与client交互计算s_test"""
for i in range(args.world_size - 1):
#id_client=random.randint(1,args.world_size) #选择client
#向选择的client发送当前v
print("send v to client:", i + 1)
for j in v:
temp = j
dist.broadcast(src=0, tensor=temp, group=group[i])
#当client计算完成，从client接收v，准备发给下一个client
print("rec v from client:", i + 1)
v_new = copy.deepcopy(v)
for j in v_new:
temp = j
dist.broadcast(src=i + 1, tensor=temp, group=group[i])
s_test.append(v_new)
#s_test计算结束，将最终s_test发送给全体client
e_s_test = s_test[0]
for i in range(1, args.world_size - 1):
e_s_test = [i + j for i, j in six.moves.zip(e_s_test, s_test[i])]
for j in e_s_test:
temp = j
dist.broadcast(src=0, tensor=temp, group=allgroup)
"""交互结束"""
#从client接收influence
print("rec influence")
allinfluence = []
influence = torch.tensor([i for i in range(4285)], dtype=torch.float32)
influence = influence.to(args.device)
for i in range(args.world_size - 1):
dist.broadcast(src=i + 1, tensor=influence, group=group[i])
temp = copy.deepcopy(influence)
allinfluence.append(temp)
torch.save(allinfluence, 'influence')
else:
"""clents"""
print("{}号服务器的第{}号GPU作为第{}个client".format(args.rank, gpu, newrank))
#设置gpu
args.device = torch.device(
'cuda:{}'.format(args.gpu)
if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
#加载训练数据
data = torch.load("data/distributed/data_of_client{}".format(newrank))
bsz = 64
train_set = torch.utils.data.DataLoader(data, batch_size=bsz)
model = CNNMnist().to(args.device)
model.load_state_dict(torch.load('w_wag')) #加载模型
data, target = train_set.dataset[0]
data = train_set.collate_fn([data])
target = train_set.collate_fn([target])
grad_v = grad_z(data, target, model, gpu=gpu)
v = grad_v
"""calculate influence function"""
v_new = []
#从server接收v
""" 和server交互计算s_test，可以循环迭代(当前只进行了一次迭代，没有循环）"""
for i in v:
temp = i
dist.broadcast(src=0, tensor=temp, group=group[newrank - 1])
v_new.append(temp)
s_test = stest(v_new,
model,
train_set,
gpu,
damp=0.01,
scale=1000.0,
repeat=5) #计算s_test
#向server发送s_test,进行下一次迭代
for i in s_test:
temp = copy.copy(i)
dist.broadcast(src=newrank, tensor=temp, group=group[newrank - 1])
#迭代完成后，从server接收最终的s_test，计算influence function
s_test_fin = []
for i in s_test:
temp = copy.copy(i)
dist.broadcast(src=0, tensor=temp, group=allgroup)
s_test_fin.append(temp)
"""s_test计算结束，得到最终的s_test_fin，开始计算influence"""
print("client:", newrank, "calculate influence")
n = len(train_set.dataset)
influence = np.array([i for i in range(n)], dtype='float32')
for i in utility.create_progressbar(len(train_set.dataset),
desc='influence',
start=0):
#计算grad
data, target = train_set.dataset[i]
data = train_set.collate_fn([data])
target = train_set.collate_fn([target])
grad_z_vec = grad_z(data, target, model, gpu=gpu)
#计算influence
inf_tmp = -sum([
torch.sum(k * j).data.cpu().numpy()
for k, j in six.moves.zip(grad_z_vec, s_test_fin)
]) / n
influence[i] = inf_tmp
influence = torch.tensor(influence).to(args.device)
#向服务器发送influence
print("client:", newrank, "send influence to server")
dist.broadcast(src=newrank, tensor=influence, group=group[newrank - 1])
print("client:", newrank, "end send influence to server")
def main_worker(gpu, ngpus_per_node, args):
print("gpu:", gpu)
args.gpu = gpu
if args.rank == 0: # (第一台服务器只有三台GPU，需要特殊处理)
newrank = args.rank * ngpus_per_node + gpu
else:
newrank = args.rank * ngpus_per_node + gpu-1
# 初始化,使用tcp方式进行通信
dist.init_process_group(init_method=args.init_method, backend="nccl", world_size=args.world_size, rank=newrank)
# 建立通信group,rank=0作为server，用broadcast模拟send和rec，需要server和每个client建立group
group = []
for i in range(1, args.world_size):
group.append(dist.new_group([0, i]))
allgroup = dist.new_group([i for i in range(args.world_size)])
if newrank == 0:
""" server"""
print("{}号服务器的第{}块GPU作为server".format(args.rank, gpu))
# 在模型训练期间，server只负责整合参数并分发，不参与任何计算
# 设置cpu
args.device = torch.device(
'cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
# 加载测试数据
trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
dataset_test = datasets.MNIST('data/', train=False, download=True, transform=trans_mnist)
test_set = torch.utils.data.DataLoader(dataset_test, batch_size=args.bs)
"""calculate influence function"""
model = CNNMnist().to(args.device)
model.load_state_dict(torch.load('w_wag'))
test_id = 0 # 选择的test数据id
data, target = test_set.dataset[test_id]
data = test_set.collate_fn([data])
target = test_set.collate_fn([target])
print("begin grad")
grad_test = grad_z(data, target, model, gpu, create_graph=False) # grad_test
print("end grad")
v = grad_test
"""server与client交互计算s_test（采用rka算法）"""
#计算模型总参数
num_parameters=0
for i in list(model.parameters()):
# 首先求出每个tensor中所含参数的个数
temp = 1
for j in i.size():
temp *= j
num_parameters+=temp
# 向client发送grad_test
for i in range(args.world_size - 1):
print("send grad_test to client:", i+1)
for j in v:
temp = j
dist.broadcast(src=0, tensor=temp, group=group[i])
for k in range(args.num_sample_rka):
#向client发送采样id
id=torch.tensor(random.randint(0,num_parameters-1)).to(args.device)
for i in range(args.world_size-1):
dist.broadcast(src=0,tensor=id,group=group[i])
# 从server接收二阶导
sec_grad = []
second_grad = [torch.zeros(list(model.parameters())[i].size()).to(args.device) for i in
range(len(list(model.parameters())))]
for i in range(args.world_size - 1):
temp = copy.deepcopy(second_grad)
for j in temp:
dist.broadcast(src=i + 1, tensor=j, group=group[i])
sec_grad.append(temp)
# 整合二阶导，然后分发给client
e_second_grad = sec_grad[0]
for i in range(1, args.world_size - 1):
e_second_grad = [i + j for i, j in six.moves.zip(e_second_grad, sec_grad[i])]
e_second_grad = [i / (args.world_size - 1) for i in e_second_grad]
for j in e_second_grad:
temp = j
dist.broadcast(src=0, tensor=temp, group=allgroup)
"""交互结束"""
# 从client接收influence
print("rec influence")
allinfluence = []
influence = torch.tensor([i for i in range(4285)], dtype=torch.float32)
influence = influence.to(args.device)
for i in range(args.world_size - 1):
dist.broadcast(src=i + 1, tensor=influence, group=group[i])
temp = copy.deepcopy(influence)
allinfluence.append(temp)
torch.save(allinfluence, 'influence/influence')
else:
"""clents"""
print("{}号服务器的第{}号GPU作为第{}个client".format(args.rank, gpu, newrank))
# 设置gpu
args.device = torch.device(
'cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
# 加载训练数据
data = torch.load("data/distributedData/data_of_client{}".format(newrank))
bsz = 64
train_set = torch.utils.data.DataLoader(data, batch_size=bsz)
model = CNNMnist().to(args.device)
model.load_state_dict(torch.load('w_wag')) # 加载模型
data, target = train_set.dataset[0]
data = train_set.collate_fn([data])
target = train_set.collate_fn([target])
grad_v = grad_z(data, target, model, gpu=gpu,create_graph=False)
grad_test = copy.deepcopy(grad_v)
"""calculate influence function"""
""" 和server交互计算s_test，可以循环迭代(采用rka算法）"""
# 从server接收grad_test
for i in grad_test:
dist.broadcast(src=0, tensor=i, group=group[newrank - 1])
stest = copy.deepcopy(grad_test)
for k in range(args.num_sample_rka):
#从server接收采样id,计算二阶导
id=torch.tensor([0]).to(args.device).to(args.device)
dist.broadcast(src=0,tensor=id,group=group[newrank - 1])
idt= id.item()
second_grad=hessian(model,train_set,idt,gpu=args.gpu)
#向server发送二阶导
for i in second_grad:
temp = i
dist.broadcast(src=newrank, tensor=temp, group=group[newrank - 1])
# 从server接收最终的二阶导
for i in second_grad:
temp = i
dist.broadcast(src=0, tensor=temp, group=allgroup)
# 使用rka算法计算stest
stest = rka(stest, second_grad, grad_test)
s_test_fin = stest
""""s_test计算结束，得到最终的s_test_fin，开始计算influence"""
print("client:", newrank, "calculate influence")
n = len(train_set.dataset)
influence = np.array([i for i in range(n)], dtype='float32')
for i in utility.create_progressbar(len(train_set.dataset), desc='influence', start=0):
# 计算grad
data, target = train_set.dataset[i]
data = train_set.collate_fn([data])
target = train_set.collate_fn([target])
grad_z_vec = grad_z(data, target, model, gpu=gpu)
# 计算influence
inf_tmp = -sum(
[torch.sum(k * j).data.cpu().numpy() for k, j in six.moves.zip(grad_z_vec, s_test_fin)]) / n
influence[i] = inf_tmp
influence = torch.tensor(influence).to(args.device)
# 向服务器发送influence
print("client:", newrank, "send influence to server")
dist.broadcast(src=newrank, tensor=influence, group=group[newrank - 1])
print("client:", newrank, "end send influence to server")
loss_locals_30.append(copy.deepcopy(loss30))
# update global weights
w_glob = FedAvg(w_locals)
w_glob1 = FedAvg(w_locals_1)
w_glob5 = FedAvg(w_locals_5)
w_glob10 = FedAvg(w_locals_10)
w_glob15 = FedAvg(w_locals_15)
w_glob20 = FedAvg(w_locals_20)
w_glob25 = FedAvg(w_locals_25)
w_glob30 = FedAvg(w_locals_30)
# copy weight to net_glob
net_glob.load_state_dict(w_glob)
net_glob1.load_state_dict(w_glob1)
net_glob5.load_state_dict(w_glob5)
net_glob10.load_state_dict(w_glob10)
net_glob15.load_state_dict(w_glob15)
net_glob20.load_state_dict(w_glob20)
net_glob25.load_state_dict(w_glob25)
net_glob30.load_state_dict(w_glob30)
# print loss
loss_avg = sum(loss_locals) / len(loss_locals)
loss_avg_1 = sum(loss_locals_1) / len(loss_locals_1)
loss_avg_5 = sum(loss_locals_5) / len(loss_locals_5)
loss_avg_10 = sum(loss_locals_10) / len(loss_locals_10)
loss_avg_15 = sum(loss_locals_15) / len(loss_locals_15)
loss_avg_20 = sum(loss_locals_20) / len(loss_locals_20)
loss_avg_25 = sum(loss_locals_25) / len(loss_locals_25)
class Client():
def __init__(self,
args,
dataset=None,
idxs=None,
w=None,
C=0.5,
sigma=0.05):
self.args = args
self.loss_func = nn.CrossEntropyLoss()
self.ldr_train = DataLoader(DatasetSplit(dataset, idxs),
batch_size=self.args.local_bs,
shuffle=True)
self.model = CNNMnist(args=args).to(args.device)
self.model.load_state_dict(w)
self.C = C
self.sigma = sigma
if self.args.mode == 'Paillier':
self.pub = pub
self.priv = priv
def train(self):
w_old = copy.deepcopy(self.model.state_dict())
net = copy.deepcopy(self.model)
net.train()
#train and update
optimizer = torch.optim.SGD(net.parameters(),
lr=self.args.lr,
momentum=self.args.momentum)
for iter in range(self.args.local_ep):
batch_loss = []
for batch_idx, (images, labels) in enumerate(self.ldr_train):
images, labels = images.to(self.args.device), labels.to(
self.args.device)
net.zero_grad()
log_probs = net(images)
loss = self.loss_func(log_probs, labels)
loss.backward()
optimizer.step()
batch_loss.append(loss.item())
w_new = net.state_dict()
update_w = {}
if self.args.mode == 'plain':
for k in w_new.keys():
update_w[k] = w_new[k] - w_old[k]
# 1. part one
# DP mechanism
elif self.args.mode == 'DP':
for k in w_new.keys():
# calculate update_w
update_w[k] = w_new[k] - w_old[k]
# clip the update
update_w[k] = update_w[k] / max(
1,
torch.norm(update_w[k], 2) / self.C)
# add noise ,cause update_w might reveal user's data also ,we should add noise before send to server
update_w[k] += np.random.normal(0.0, self.sigma**2 * self.C**2)
# 2. part two
# Paillier enc
elif self.args.mode == 'Paillier':
print(len(w_new.keys()))
for k in w_new.keys():
print("start ", k, flush=True)
update_w[k] = w_new[k] - w_old[k]
update_w_list = update_w[k].view(-1).cpu().tolist()
for iter, w in enumerate(update_w_list):
update_w_list[iter] = self.pub.encrypt(w)
update_w[k] = update_w_list
print("end ", flush=True)
else:
exit()
return update_w, sum(batch_loss) / len(batch_loss)
def update(self, w_glob):
if self.args.mode == 'plain':
self.model.load_state_dict(w_glob)
elif self.args.mode == 'DP':
self.model.load_state_dict(w_glob)
elif self.args.mode == 'Paillier':
w_glob_ciph = copy.deepcopy(w_glob)
for k in w_glob_ciph.keys():
for iter, item in enumerate(w_glob_ciph[k]):
w_glob_ciph[k][iter] = self.priv.decrypt(item)
shape = list(self.model.state_dict()[k].size())
w_glob_ciph[k] = torch.FloatTensor(w_glob_ciph[k]).to(
self.args.device).view(*shape)
self.model.state_dict()[k] += w_glob_ciph[k]
else:
exit()
#5 MALICIOUS
w_locals_5.append(copy.deepcopy(w5))
loss_locals_5.append(copy.deepcopy(loss5))
#10 MALICIOUS
w_locals_10.append(copy.deepcopy(w10))
loss_locals_10.append(copy.deepcopy(loss10))
# update global weights
w_glob = FedAvg(w_locals)
w_glob_5 = FedAvg(w_locals_5)
w_glob_10 = FedAvg(w_locals_10)
# copy weight to net_glob
net_glob.load_state_dict(w_glob)
net_glob5.load_state_dict(w_glob_5)
net_glob10.load_state_dict(w_glob_10)
# print loss
loss_avg = sum(loss_locals) / len(loss_locals)
loss_avg_5 = sum(loss_locals_5) / len(loss_locals_5)
loss_avg_10 = sum(loss_locals_10) / len(loss_locals_10)
non_malicious_structure[iter][0] = loss_avg
non_malicious_structure5[iter][0] = loss_avg_5
non_malicious_structure10[iter][0] = loss_avg_10
malicious_structure5[iter][0] = loss_avg_5
malicious_structure10[iter][0] = loss_avg_10
print('NO ATTACK ---> Round {:3d}, Average loss {:.3f}'.format(
class Server():
def __init__(self, args, w):
self.args = args
self.clients_update_w = []
self.clients_loss = []
self.model = CNNMnist(args=args).to(args.device)
self.model.load_state_dict(w)
def FedAvg(self):
# 1. part one
# DP mechanism
# cause we choose to add noise at client end,the fedavg should be the same as plain
if self.args.mode == 'plain' or self.args.mode == 'DP':
update_w_avg = copy.deepcopy(self.clients_update_w[0])
for k in update_w_avg.keys():
for i in range(1, len(self.clients_update_w)):
update_w_avg[k] += self.clients_update_w[i][k]
update_w_avg[k] = torch.div(update_w_avg[k],
len(self.clients_update_w))
self.model.state_dict()[k] += update_w_avg[k]
return copy.deepcopy(self.model.state_dict()), sum(
self.clients_loss) / len(self.clients_loss)
# 2. part two
# Paillier add
elif self.args.mode == 'Paillier':
update_w_avg = copy.deepcopy(self.clients_update_w[0])
for k in update_w_avg.keys():
client_num = len(self.clients_update_w)
for i in range(1, client_num):
for iter in range(len(update_w_avg[k])):
update_w_avg[k][iter] += self.clients_update_w[i][k][
iter]
for iter in range(len(update_w_avg[k])):
update_w_avg[k][iter] /= client_num
return update_w_avg, sum(self.clients_loss) / len(
self.clients_loss)
else:
exit()
def test(self, datatest):
self.model.eval()
# testing
test_loss = 0
correct = 0
data_loader = DataLoader(datatest, batch_size=self.args.bs)
for idx, (data, target) in enumerate(data_loader):
if self.args.gpu != -1:
data, target = data.cuda(), target.cuda()
log_probs = self.model(data)
# sum up batch loss
test_loss += F.cross_entropy(log_probs, target,
reduction='sum').item()
# get the index of the max log-probability
y_pred = log_probs.data.max(1, keepdim=True)[1]
correct += y_pred.eq(
target.data.view_as(y_pred)).long().cpu().sum()
test_loss /= len(data_loader.dataset)
accuracy = 100.00 * correct / len(data_loader.dataset)
return accuracy, test_loss