### 联邦平均算法简介 联邦平均（FedAvg）是一种用于分布式机器学习的技术，特别适用于保护隐私的数据分布环境。该方法允许多个客户端在本地训练模型，并仅共享更新后的参数给服务器端汇总，从而减少数据传输并增强安全性[^1]。 ### 使用 PyTorch 实现 FedAvg 的基本框架 为了构建一个简单的基于 PyTorch 的 FedAvg 流程，可以按照如下方式组织代码结构： #### 初始化全局模型和服务端设置 ```python import torch.nn as nn from torchvision import datasets, transforms from torch.utils.data import DataLoader class Net(nn.Module): def __init__(self): super(Net, self).__init__() # 定义网络层... def forward(self, x): # 前向传播逻辑... def init_global_model(): global_model = Net() return global_model.cuda() if use_cuda else global_model ``` #### 数据加载与分配至各客户端 考虑到不同设备上的异构特性，需合理划分数据集到各个参与方手中： ```python transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform) # 将整个训练集划分为 n_clients 份子集 client_datasets = [] for i in range(n_clients): start_idx = int(i * len(train_dataset)/n_clients) end_idx = min(int((i+1)*len(train_dataset)/n_clients), len(train_dataset)) client_datasets.append(DataLoader( dataset=train_dataset[start_idx:end_idx], batch_size=batch_size, shuffle=True)) test_loader = DataLoader( datasets.MNIST('./data', train=False, transform=transform), batch_size=test_batch_size, shuffle=False) ``` #### 客户端局部更新过程 每个客户机将在自己的私有数据上执行一轮或多轮SGD迭代来优化其副本权重: ```python def local_update(client_id, model, optimizer, loss_fn, epochs=EPOCHS_PER_ROUND): device = 'cuda' if use_cuda and torch.cuda.is_available() else 'cpu' for epoch in range(epochs): running_loss = 0.0 correct_predictions = 0 for data, target in client_datasets[client_id]: data, target = data.to(device), target.to(device) output = model(data) loss = loss_fn(output, target) _, preds = torch.max(output, dim=1) correct_predictions += torch.sum(preds == target).item() optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item()*data.size(0) avg_accuracy = float(correct_predictions / len(client_datasets[client_id].dataset)) print(f'[Client {client_id}] Epoch [{epoch}/{EPOCHS_PER_ROUND}], Loss: {running_loss/len(client_datasets[client_id]):.4f}, Accuracy: {avg_accuracy:.4f}') return model.state_dict(), avg_accuracy ``` #### 参数聚合机制 服务端收集来自所有选定客户的梯度变化量，并计算加权均值得到新的全局状态字典： ```python def aggregate_weights(weighted_updates): aggregated_state_dict = {} total_weight = sum([w for w,_ in weighted_updates]) for key in global_model.state_dict().keys(): averaged_param = None for weight, update in weighted_updates: param = update[key]*weight if averaged_param is None: averaged_param = param.clone() else: averaged_param.add_(param) aggregated_state_dict[key] = averaged_param.div(total_weight) return aggregated_state_dict ``` #### 主循环控制通信回合数 最后，在主程序里定义好必要的超参配置后启动多轮次的协同工作流： ```python if __name__ == '__main__': num_rounds = NUM_ROUNDS clients_per_round = CLIENTS_PER_ROUND for round_num in range(num_rounds): selected_client_ids = np.random.choice(range(n_clients), size=min(clients_per_round,n_clients), replace=False) updates = [] for cid in selected_client_ids: updated_params, acc = local_update(cid, copy.deepcopy(global_model), optimizers[cid], criterion) updates.append((acc,updated_params)) new_global_state_dict = aggregate_weights(updates) global_model.load_state_dict(new_global_state_dict) test_acc = evaluate_on_test_set(global_model) print(f'\nRound #{round_num}: Test Set Accuracy={test_acc}\n') ``` 上述代码片段展示了如何利用 PyTorch 构建一套简易版的 FedAvg 系统原型[^3]。