class Client:
def __init__(self, name, train_data_dir, test_data_dir):
self.name = name
transform = transforms.ToTensor()
trainset = datasets.ImageFolder(train_data_dir, transform=transform)
self.trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=BATCH_SIZE,
shuffle=True
)
testset = datasets.ImageFolder(test_data_dir, transform=transform)
self.testloader = torch.utils.data.DataLoader(
testset,
batch_size=BATCH_SIZE,
shuffle=False
)
dataset_list = list(self.trainloader)
self.dataset_len = len(dataset_list)
self.net = LeNet().to(device)
self.criterion = nn.CrossEntropyLoss()
def update(self, net_dict, center_params_dict):
self.net.load_state_dict(net_dict)
for i in range(LOCAL_EPOCH_NUM):
data_iter = iter(self.trainloader)
for b in range(self.dataset_len):
inputs, labels = next(data_iter)
inputs = torch.index_select(inputs, 1, torch.LongTensor([0]))
inputs, labels = inputs.to(device), labels.to(device)
outputs = self.net(inputs)
loss = self.criterion(outputs, labels)
optimizer = optim.SGD(self.net.parameters(), lr=LR, momentum=0.9)
optimizer.zero_grad()
loss.backward()
params_modules = list(self.net.named_parameters())
for params_module in params_modules:
name, params = params_module
params.grad += MU * (params.data - center_params_dict[name])
optimizer.step()
return self.net.state_dict()
class Client:
def __init__(self, name, train_data_dir, test_data_dir, pk, sk):
self.name = name
self.pk = pk
self.sk = sk
transform = transforms.ToTensor()
trainset = datasets.ImageFolder(train_data_dir, transform=transform)
self.trainloader = torch.utils.data.DataLoader(trainset,
batch_size=BATCH_SIZE,
shuffle=True)
testset = datasets.ImageFolder(test_data_dir, transform=transform)
self.testloader = torch.utils.data.DataLoader(testset,
batch_size=BATCH_SIZE,
shuffle=False)
dataset_list = list(self.trainloader)
self.dataset_len = len(dataset_list)
self.net = LeNet().to(device)
self.criterion = nn.CrossEntropyLoss()
def get_encrypted_grad(self, client_inputs, client_labels, net_dict):
self.net.load_state_dict(net_dict)
client_outputs = self.net(client_inputs)
client_loss = self.criterion(client_outputs, client_labels)
client_optimizer = optim.SGD(self.net.parameters(),
lr=LR,
momentum=0.9)
client_optimizer.zero_grad()
client_loss.backward()
params_modules = list(self.net.named_parameters())
params_grad_list = []
for params_module in params_modules:
name, params = params_module
params_grad_list.append(copy.deepcopy(params.grad).view(-1))
params_grad = ((torch.cat(params_grad_list, 0) + bound) *
2**prec).long().cuda()
client_encrypted_grad = Enc(self.pk, params_grad)
client_optimizer.zero_grad()
return client_encrypted_grad
class Client:
def __init__(self, name, train_data_dir, test_data_dir):
self.name = name
transform = transforms.ToTensor()
trainset = datasets.ImageFolder(train_data_dir, transform=transform)
self.trainloader = torch.utils.data.DataLoader(trainset,
batch_size=BATCH_SIZE,
shuffle=True)
testset = datasets.ImageFolder(test_data_dir, transform=transform)
self.testloader = torch.utils.data.DataLoader(testset,
batch_size=BATCH_SIZE,
shuffle=False)
dataset_list = list(self.trainloader)
self.dataset_len = len(dataset_list)
self.net = LeNet().to(device)
self.criterion = nn.CrossEntropyLoss()
def get_grad(self, client_inputs, client_labels, net_dict):
self.net.load_state_dict(net_dict)
client_outputs = self.net(client_inputs)
client_loss = self.criterion(client_outputs, client_labels)
client_optimizer = optim.SGD(self.net.parameters(),
lr=LR,
momentum=0.9)
client_optimizer.zero_grad()
client_loss.backward()
client_grad_dict = dict()
params_modules = list(self.net.named_parameters())
for params_module in params_modules:
name, params = params_module
params_grad = copy.deepcopy(params.grad)
client_grad_dict[name] = params_grad
client_optimizer.zero_grad()
return client_grad_dict
], net_dict, attacker_net)
# client_attacker_grad_dict = get_client_grad_model(client_attacker_inputs, client_attacker_labels, attack_module)
# 取各client参数梯度均值
client_average_grad_dict = dict()
# attacker_average_grad_dict = dict()
# for key in client_attacker_grad_dict:
# attacker_average_grad_dict[key] = client_attacker_grad_dict[key]
for key in client_0_grad_dict:
client_average_grad_dict[key] = client_0_grad_dict[key] * (
1 / Num_client) + client_1_grad_dict[key] * (
1 / Num_client) + client_2_grad_dict[key] * (
1 / Num_client) + client_attacker_grad_dict[key] * (
1 / Num_client)
# 加载梯度
params_modules_server = net.named_parameters()
# params_modules_attacker = attacker_net.named_parameters()
# for params_module in params_modules_attacker:
# (name_attacker, params) = params_module
# params.grad = attacker_average_grad_dict[name_attacker]
# optimizer_backdoor.step()
for params_module in params_modules_server:
(name, params) = params_module
params.grad = client_average_grad_dict[
name] # 用字典中存储的子模型的梯度覆盖server中的参数梯度
optimizer_server.step()
# 每跑完一次epoch测试一下准确率
with torch.no_grad():
correct = 0
total = 0
client_attacker_net)
# 取各client参数梯度均值
client_average_grad_dict = dict()
attacker_average_grad_dict = dict()
for key in client_attacker_grad_dict:
attacker_average_grad_dict[key] = client_attacker_grad_dict[
key]
for key in client_0_grad_dict:
client_average_grad_dict[key] = client_0_grad_dict[key] * (
1 / Num_client) + client_1_grad_dict[key] * (
1 / Num_client) + client_2_grad_dict[key] * (
1 / Num_client
) + client_attacker_grad_dict[key] * (1 / Num_client)
# 加载梯度
params_modules_server = net.named_parameters()
params_modules_attacker = attacker_net.named_parameters()
for params_module in params_modules_attacker:
(name_attacker, params) = params_module
params.grad = client_average_grad_dict[name_attacker]
optimizer_backdoor.step()
for params_module in params_modules_server:
(name, params) = params_module
params.grad = client_average_grad_dict[
name] # 用字典中存储的子模型的梯度覆盖server中的参数梯度
optimizer_server.step()
# # 每训练100个batch打印一次平均loss
# sum_loss += loss_c0.item()
# if i % 100 == 99:
# print('[%d, %d] loss: %.03f'
# % (epoch + 1, i + 1, sum_loss / 100))
