第二轮联邦学习的全局模型和第一轮联邦学习的全局模型保持一致的代码

好的，以下是一个简单的代码示例，用于在第二轮联邦学习中将全局模型设为第一轮联邦学习的全局模型：

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch import nn, optim

# 定义模型结构
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self.dropout = nn.Dropout(p=0.2)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.dropout(self.relu(self.fc2(x)))
        x = self.fc3(x)
        return x

# 定义联邦学习服务器
class FederatedServer:
    def __init__(self, num_clients, train_data, test_data, lr=0.01, batch_size=64, epochs=10):
        self.num_clients = num_clients
        self.train_data = train_data
        self.test_data = test_data
        self.lr = lr
        self.batch_size = batch_size
        self.epochs = epochs
        self.clients = []
        self.server_model = Model()
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.SGD(self.server_model.parameters(), lr=self.lr)

        # 初始化客户端
        for i in range(num_clients):
            data_loader = DataLoader(train_data[i], batch_size=self.batch_size, shuffle=True)
            client_model = Model()
            client_optimizer = optim.SGD(client_model.parameters(), lr=self.lr)
            self.clients.append({'data_loader': data_loader, 'model': client_model, 'optimizer': client_optimizer})

    # 训练客户端模型
    def train_client_model(self, client):
        client['model'].train()
        for epoch in range(self.epochs):
            for images, labels in client['data_loader']:
                client['optimizer'].zero_grad()
                output = client['model'](images)
                loss = self.criterion(output, labels)
                loss.backward()
                client['optimizer'].step()

    # 聚合客户端模型
    def aggregate_client_models(self):
        for param, global_param in zip(self.server_model.parameters(), self.global_model.parameters()):
            param.data = global_param.data

        for client in self.clients:
            for param, client_param in zip(self.server_model.parameters(), client['model'].parameters()):
                param.data += client_param.data / self.num_clients

    # 在测试集上评估模型
    def evaluate_model(self):
        self.server_model.eval()
        test_loss = 0
        test_accuracy = 0
        with torch.no_grad():
            for images, labels in self.test_data:
                output = self.server_model(images)
                test_loss += self.criterion(output, labels)
                ps = torch.exp(output)
                top_p, top_class = ps.topk(1, dim=1)
                equals = top_class == labels.view(*top_class.shape)
                test_accuracy += torch.mean(equals.type(torch.FloatTensor))
        return test_loss / len(self.test_data), test_accuracy / len(self.test_data)

    # 训练联邦模型
    def train(self, global_model=None):
        self.global_model = global_model
        if self.global_model is not None:
            self.server_model.load_state_dict(self.global_model.state_dict())

        for epoch in range(self.epochs):
            for client in self.clients:
                self.train_client_model(client)
            self.aggregate_client_models()
            test_loss, test_accuracy = self.evaluate_model()
            print(f"Epoch {epoch+1}/{self.epochs}, Test Loss: {test_loss:.3f}, Test Accuracy: {test_accuracy:.3f}")
        
        # 返回更新后的全局模型
        return self.server_model

这段代码与之前的联邦学习服务器代码相似，除了添加了一个名为 `global_model` 的参数。在 `train` 方法中，如果 `global_model` 不为 `None`，则会将 `global_model` 的参数加载到服务器的模型中。这样，第二轮联邦学习的全局模型就与第一轮联邦学习的全局模型保持一致了。