从FedAvg到实战：用PyTorch复现联邦学习经典论文中的MNIST实验（附完整代码）
从理论到实践：PyTorch实现联邦学习FedAvg算法的MNIST实验全解析
联邦学习（Federated Learning）作为分布式机器学习的前沿技术，正在重塑隐私保护计算的格局。本文将带您深入FedAvg算法的实现细节，通过PyTorch框架完整复现MNIST数据集上的关键实验。不同于传统集中式训练，联邦学习的核心挑战在于处理非独立同分布（Non-IID）数据，这正是现实场景中的常态。
1. 实验环境搭建与数据准备
1.1 基础环境配置
首先需要配置Python 3.8+和PyTorch 1.10+环境。推荐使用conda创建虚拟环境：
conda create -n fl_env python=3.8 conda activate fl_env pip install torch torchvision matplotlib1.2 MNIST数据集的联邦化处理
传统MNIST数据集包含6万张手写数字图像，我们需要将其模拟为分布式场景：
from torchvision import datasets, transforms from torch.utils.data import DataLoader, Subset import numpy as np def prepare_federated_datasets(num_clients=100, iid=True): transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('./data', train=False, transform=transform) if iid: # IID划分：随机打乱后均匀分配 client_indices = np.random.permutation(len(train_dataset)) client_indices = np.array_split(client_indices, num_clients) else: # Non-IID划分：按标签排序后分片分配 sorted_indices = np.argsort(train_dataset.targets.numpy()) shards = np.array_split(sorted_indices, num_clients * 2) client_indices = [np.concatenate(shards[i::num_clients]) for i in range(num_clients)] client_datasets = [ Subset(train_dataset, indices) for indices in client_indices ] return client_datasets, test_dataset关键参数说明：
num_clients：客户端数量（默认100）iid：数据分布类型（True为IID，False为Non-IID）
2. FedAvg算法核心实现
2.1 服务器端聚合逻辑
服务器负责协调全局模型参数聚合：
import copy import torch class FedAvgServer: def __init__(self, model, clients, test_loader): self.global_model = model self.clients = clients self.test_loader = test_loader def aggregate(self, client_weights, client_sizes): total_size = sum(client_sizes) averaged_weights = {} for key in client_weights[0].keys(): averaged_weights[key] = sum( [weights[key] * size for weights, size in zip(client_weights, client_sizes)] ) / total_size self.global_model.load_state_dict(averaged_weights) def evaluate(self): self.global_model.eval() correct = 0 total = 0 with torch.no_grad(): for data, target in self.test_loader: output = self.global_model(data) pred = output.argmax(dim=1) correct += (pred == target).sum().item() total += target.size(0) return correct / total2.2 客户端本地训练
每个客户端基于本地数据更新模型：
class FedAvgClient: def __init__(self, dataset, model): self.dataset = dataset self.model = copy.deepcopy(model) self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01) self.criterion = torch.nn.CrossEntropyLoss() def train(self, epochs=1, batch_size=10): loader = DataLoader(self.dataset, batch_size=batch_size, shuffle=True) self.model.train() for _ in range(epochs): for data, target in loader: self.optimizer.zero_grad() output = self.model(data) loss = self.criterion(output, target) loss.backward() self.optimizer.step() return self.model.state_dict()3. 关键超参数实验分析
3.1 通信轮次与模型精度关系
我们通过控制变量法研究三个核心参数的影响：
| 参数 | 含义 | 典型取值 |
|---|---|---|
| C | 每轮参与的客户端比例 | 0.1-1.0 |
| E | 本地训练epoch数 | 1-20 |
| B | 本地batch大小 | 10-600 |
实验结果对比：
def run_experiment(model, clients, test_loader, C=0.1, E=5, B=10, rounds=20): server = FedAvgServer(model, clients, test_loader) accuracies = [] for r in range(rounds): selected_clients = np.random.choice( clients, size=max(1, int(C * len(clients))), replace=False ) client_weights = [] client_sizes = [] for client in selected_clients: weights = client.train(epochs=E, batch_size=B) client_weights.append(weights) client_sizes.append(len(client.dataset)) server.aggregate(client_weights, client_sizes) acc = server.evaluate() accuracies.append(acc) return accuracies3.2 Non-IID场景下的挑战
非独立同分布数据会显著影响模型收敛速度。实验表明：
- IID数据：通常100轮内达到95%+准确率
- Non-IID数据：需要200-300轮才能达到相似水平
提示：在Non-IID场景下，建议增加本地训练轮次(E)并减小batch大小(B)，这有助于客户端更好地学习本地数据特征
4. 完整实验流程与结果可视化
4.1 端到端实验流程
# 初始化模型 class CNNModel(torch.nn.Module): def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(1, 32, 5) self.conv2 = torch.nn.Conv2d(32, 64, 5) self.fc1 = torch.nn.Linear(1024, 512) self.fc2 = torch.nn.Linear(512, 10) def forward(self, x): x = torch.relu(torch.max_pool2d(self.conv1(x), 2)) x = torch.relu(torch.max_pool2d(self.conv2(x), 2)) x = x.view(-1, 1024) x = torch.relu(self.fc1(x)) return self.fc2(x) # 准备数据 clients, test_dataset = prepare_federated_datasets(num_clients=100, iid=False) test_loader = DataLoader(test_dataset, batch_size=128) # 运行实验 model = CNNModel() clients = [FedAvgClient(ds, model) for ds in clients] accuracies = run_experiment(model, clients, test_loader, C=0.1, E=5, B=50) # 结果可视化 plt.plot(accuracies) plt.xlabel('Communication Rounds') plt.ylabel('Test Accuracy') plt.title('FedAvg on Non-IID MNIST') plt.grid() plt.show()4.2 性能优化技巧
- 动态调整学习率：随着训练轮次增加，逐步降低学习率
- 客户端选择策略：优先选择损失下降快的客户端
- 模型压缩：上传参数前进行量化或稀疏化
典型收敛曲线对比：
IID数据收敛曲线：平滑快速上升 Non-IID数据收敛曲线：波动较大，前期上升缓慢5. 扩展应用与进阶方向
5.1 扩展到其他数据集
FedAvg同样适用于CIFAR-10等更复杂数据集，但需注意：
- 模型架构需要调整（更深的CNN）
- 通信成本会显著增加
- 可能需要更多客户端参与
5.2 联邦学习的未来演进
- 个性化联邦学习：允许客户端保留部分个性化参数
- 异步联邦学习：放松同步更新要求
- 跨模态联邦学习：处理多模态数据场景
在实际项目中，我们发现当客户端数据分布差异较大时，适当增加本地训练轮次（E=5-10）能显著提升最终模型性能。同时，使用动量优化器替代普通SGD也能加速收敛过程。