用Python写一段联邦学习的代码

# 导入必要的库
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# 初始化超参数
num_clients = 10 # 客户端数量
num_epochs = 5 # 每轮训练的迭代次数
batch_size = 32 # 每个客户端的批次大小
learning_rate = 0.01 # 学习率

# 初始化模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(5, 1)

def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x

# 初始化数据
x_train = np.random.rand(100, 10)
y_train = (x_train[:,0] + x_train[:,1] + x_train[:,2] > 1.5).astype(float).reshape(-1, 1)
train_data = [(x_train[i], y_train[i]) for i in range(len(x_train))]

# 分割数据
num_data_per_client = len(train_data) // num_clients
train_data_split = [train_data[i:i+num_data_per_client] for i in range(0, len(train_data), num_data_per_client)]

# 初始化客户端
clients = []
for i in range(num_clients):
client_model = Model()
client_optimizer = torch.optim.Adam(client_model.parameters(), lr=learning_rate)
clients.append((client_model, client_optimizer))

# 训练过程
for epoch in range(num_epochs):
for client_id in range(num_clients):
# 获取客户端的模型与优化器
client_model, client_optimizer = clients[client_id]
client_data_loader = DataLoader(train_data_split[client_id], batch_size=batch_size, shuffle=True)

# 模型训练
for input, label in client_data_loader:
input, label = torch.Tensor(input), torch.Tensor(label)
client_optimizer.zero_grad()
loss = nn.BCEWithLogitsLoss()(client_model(input), label)
loss.backward()
client_optimizer.step()

# 聚合模型
if client_id == 0:
aggr_model = client_model
else:
for aggr_param, client_param in zip(aggr_model.parameters(), client_model.parameters()):
aggr_param.data += client_param.data

# 对聚合模型参数取平均值
for aggr_param in aggr_model.parameters():
aggr_param.data /= num_clients

# 测试模型
x_test = np.random.rand(10, 10)
y_test = (x_test[:,0] + x_test[:,1] + x_test[:,2] > 1.5).astype(float).reshape(-1, 1)
test_data = [(x_test[i], y_test[i]) for i in range(len(x_test))]
test_data_loader = DataLoader(test_data, batch_size=1, shuffle=False)

correct = 0
total = 0
for input, label in test_data_loader:
input, label = torch.Tensor(input), torch.Tensor(label)
output = aggr_model(input)
predicted = (output > 0.5).float()
total += label.size(0)
correct += (predicted == label).sum().item()

print('Accuracy of the model on the test data: {:.2f}%'.format(100 * correct / total))