联邦学习实战：端云图像分类在纯CPU环境下的FedAvg与Non-IID实现

几个月前我在本机 x86 环境里跑通了一个端云联邦图像分类应用，从客户端本地训练、服务端聚合，到最终的可视化对比，整个过程踩了不少坑，也把联邦学习里最容易让人困惑的几个环节彻底捋清楚了。这篇就当是给后来人的一份实操笔记。

这个项目实际解决的问题很直白：图像数据分散在本地、不能集中上传，但又要训练一个能用的图像分类模型。端云联邦的具体形态是——云侧维护一个全局模型，端侧用本地数据做训练，只上传模型权重更新，原始数据不出设备。我这次用的是 x86 平台的 PC 级硬件，没有 GPU，纯 CPU 训练，跑 CIFAR-10 这类公开数据集，完全能验证链路。

适合看这篇文章的，是你已经会用 PyTorch 写了几个分类网络，想进一步把分布式训练、隐私保护这套东西落地，但还没想清楚通信、聚合、数据切分这些细节。我会把从零搭一个联邦图像分类应用的完整链路拆开讲清楚，包括为什么这么设计、代码怎么写、跑了哪些实验、有哪些坑是我替你踩过的。

1. 先拆解这个项目在做什么：端云联邦图像分类的核心逻辑

其实端云联邦这个叫法听起来高大上，但剥开来看，它只是换了一种训练范式。集中式训练里，数据集中到一台服务器，训练出全局模型；联邦训练里，云端服务器持有全局模型，但训练不出门——客户端拿本地数据在本地算，算完把“模型更新”交回云端，云端再把各方更新融合起来。

1.1 端与云各自的分工

在这个项目里，我把端侧和云侧的责任划分得尽量清晰：

• 端侧（Client）：持有本地数据集，负责加载最新全局模型，在本地数据上跑若干轮梯度下降，产生新一轮权重，把权重（或权重差值）传回云端。
• 云侧（Server）：初始化全局模型，选择参与本轮训练的一组客户端，广播模型，接收各客户端返回的权重更新，执行聚合得到新全局模型，再进入下一轮。

这个分工和真实场景是完全一致的。你可以想到手机、门禁设备、工业相机，它们就是端侧；承载聚合计算的那台服务器就是云侧。我在 x86 上做的是模拟，但架构上没有缩水。

代码层面，我特意把 server 和 client 拆成两个独立模块，中间用一层通信抽象隔离开。这样以后从“单机模拟”切到“多机部署”，只需要替换通信实现，训练逻辑完全不用动。这也是整个项目里我觉得最有价值的架构决策。

1.2 为什么聚合策略选 FedAvg 而不是别的

联邦学习里的核心算法是服务端怎么融合各客户端上传的模型。最常见的基线算法是 FedAvg——联邦平均。它的思想比听起来还简单：把各客户端上传的权重，按各客户端本地样本数加权平均。

[ w_{global} = \sum_{k=1}^{K} \frac{n_k}{n} w_k ]

其中 (n_k) 是第 (k) 个客户端的样本数，(n) 是所有参与客户端样本总数，(w_k) 是第 (k) 个客户端本地训练后的权重。

选它有两个硬理由。第一，实现门槛低，不依赖额外通信轮次，对 CPU 环境友好；第二，它已经是这个领域的事实标准，后续想换 FedProx、SCAFFOLD 这些进阶算法，底层的通信和训练框架不用动，只改聚合函数。我先把 FedAvg 跑通，再谈优化，是性价比最高的路线。

1.3 为什么从 x86 环境切入

可能有人会问，联邦学习不是应该跑在手机或者边缘设备上吗，为什么用 x86？

我的实践感受是：x86 环境最大的价值在于可调试性和可复现性。移动端部署还要考虑交叉编译、系统权限、电源策略，这些都会干扰算法验证。而 x86 PC 上装好 Python 和 PyTorch，十分钟就能把联邦训练的模拟链路跑起来。等算法验证完毕，再考虑把客户端逻辑移植到真实端侧，思路会清晰很多。

而且很多实际落地场景里，云侧和端侧并不是一边是服务器一边是手机，而是多台 x86 工控机之间做协同，这个时候 x86 就是生产环境本身。所以用 x86 做实验环境，既不丢通用性，又能快速迭代。

2. x86 环境准备：依赖安装、数据切分与性能调优

先把环境准备好。我用的是一台普通的 x86 台式机，16GB 内存，CPU 是六核十二线程，没有独立 GPU。操作系统是 Ubuntu，PyTorch 用的 CPU 版本。Windows 也能跑，但是多进程行为有点区别，后面我会专门提到。

2.1 环境与目录规划

我建了下面这个目录结构，改起来很顺手：

fed_image_classifier/ ├── config.py # 全局配置：路径、超参、客户端数量 ├── data_partition.py # 数据加载与 IID / Non-IID 切分 ├── model.py # 图像分类网络定义 ├── client.py # 端侧训练逻辑 ├── server.py # 云侧聚合主循环 ├── metrics.py # 评估与训练曲线记录 ├── run_experiments.py # 一键跑多组实验 └── proto/ └── federated.proto # gRPC 通信协议定义（进阶用）

这个结构的好处是每个模块只干一件事。model.py 只管网络结构；client.py 只管本地训练；server.py 只管调度和聚合。你后期改任何一块，都不影响其他模块。很多初学者喜欢把所有代码堆在一个文件里，跑通了还好，一旦要调 Non-IID 数据分布或者加通信异常处理，就会陷入改一行崩三行的痛苦。

依赖安装我用了一个干净的虚拟环境：

python3 -m venv .venv source .venv/bin/activate pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu pip install matplotlib grpcio grpcio-tools

强烈建议先装 CPU 版 PyTorch。GPU 版在这个项目里一是没必要，二是会引入 CUDA 版本匹配问题，纯属给自己添堵。

2.2 CIFAR-10 数据集准备与切分策略

图像分类任务我直接用了 CIFAR-10 公开数据集，一共 6 万张 32×32 的彩色图片，10 个类别。集中在单机上做联邦学习研究，CIFAR-10 是最常用的基准，因为足够小、跑得快，而且类别均衡，方便观察不同数据分布的影响。

端云联邦实验里，数据切分是决定实验结果走向的关键环节。我实现了两种切分方式：

• IID（独立同分布）切分：把所有样本随机打乱，均匀分给 N 个客户端。每个客户端手里的数据分布和全局分布差不多。
• Non-IID（非独立同分布）切分：模拟真实场景——每个客户端只持有少数几个类别的数据。比如客户端 A 只有“猫”、“狗”，客户端 B 只有“卡车”、“飞机”。

Non-IID 切分的实现思路是按标签排序后做分片。我把每个类别的样本分成若干片，然后按“每个客户端分配固定类别”的方式重组。代码大概是：

def partition_non_iid(labels, num_clients, num_shards=2): # 每个类别下的样本索引 class_idx = {} for i, label in enumerate(labels): class_idx.setdefault(label, []).append(i) # 把每个类别的样本等分成 num_shards 份 shards = [] for label, idxs in class_idx.items(): random.shuffle(idxs) split = np.array_split(idxs, num_shards) for s in split: shards.append(s.tolist()) # 每个客户端分配 num_shards 个 shard random.shuffle(shards) client_data = [[] for _ in range(num_clients)] for c in range(num_clients): for _ in range(num_shards): if shards: client_data[c].extend(shards.pop()) return client_data

这里num_shards=2意味着每个客户端持有 2 个 shard。CIFAR-10 有 10 个类别，一个类别一个 shard，客户端分到 2 个 shard 就近似只持有 2 个类别的数据。这是经典的 Non-IID 构造方式，实验效果非常直观。

2.3 CPU 环境下容易被忽视的三个性能开关

x86 纯 CPU 环境跑联邦学习，性能瓶颈比 GPU 明显得多。我调了好几轮，发现有三个开关对效率影响极大：

第一个是线程数设置。PyTorch 的 CPU 算子默认会尝试用满所有核，但线程切换开销可能抵消计算收益。我在每个进程入口设置了：

import torch torch.set_num_threads(6)

如果客户端进程开得比较多，比如同时开 10 个客户端进程，每个都抢 CPU 资源，总吞吐量反而会下降。我的经验是：客户端数量乘以单客户端线程数，不要超过物理核心数的两倍。

第二个是 DataLoader 的 num_workers。CPU 训练时数据加载瓶颈非常明显。我最初没设 num_workers，每个 epoch 的数据加载时间几乎和计算时间相当。后来设为 4，训练速度直接提了约 30%。但注意 Windows 上如果 num_workers 大于 0，数据加载代码必须包在if __name__ == "__main__"里，否则会无限开进程。

第三个是推理批次大小。每轮全局模型下发后，服务端要用验证集评估准确率，这个评估过程 batch size 最好调大一些（比如 512），能明显缩短验证时间；而客户端本地训练 batch size 则保持 64 左右，避免太小导致梯度震荡。

3. 核心代码实现：从本地训练到全局聚合的一条完整链路

这一节是全文的核心。我直接给出可运行的关键代码，并解释每一段到底在干什么。先把模型写了。

3.1 图像分类模型设计

CIFAR-10 分辨率为 32×32，不需要太深的网络。我用了一个三层卷积加两层全连接的小网络，参数量约 120 万，CPU 上一个 batch 的推理时间在毫秒级，非常适合做联邦学习实验：

import torch import torch.nn as nn class FedCNN(nn.Module): def __init__(self): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(128 * 4 * 4, 256), nn.ReLU(), nn.Linear(256, 10), ) def forward(self, x): return self.classifier(self.features(x))

为什么这么设计？三个卷积层足够提取 32×32 图像的中层特征，再加两个全连接层做分类映射。参数量少，通信开销就小——联邦学习里每次通信传输的都是整个模型权重，模型越大，带宽成本越高。小网络也让单机模拟的迭代速度快很多。

3.2 客户端本地训练逻辑

端侧的训练流程本质是标准的监督分类训练，但有两个关键细节：必须从全局权重开始、只训练少量 epoch。

def local_train(model, global_state, train_loader, epochs=3, lr=0.01): # 将全局模型参数加载到本地模型 model.load_state_dict(global_state) optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) criterion = nn.CrossEntropyLoss() model.train() for epoch in range(epochs): for images, labels in train_loader: optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 返回训练后的完整权重，而不是梯度 return {k: v.clone() for k, v in model.state_dict().items()}

注意几个细节。第一，每个客户端本地训练只跑 3 到 5 个 epoch。联邦学习里本地训练太久会导致“模型漂移”，即各个客户端的模型越走越远，全局聚合后效果反而变差。第二，返回的是权重而非梯度，这样服务端聚合逻辑更简单直接。第三，使用 momentum=0.9 的 SGD，这在联邦学习场景下比 Adam 更常见，因为 Adam 的动量信息是每个客户端独立维护的，聚合后匹配不上，效果会打折扣。

3.3 服务端 FedAvg 聚合主循环

服务端做的事情可以抽象成四步：选客户端、发模型、收更新、加权平均。一次训练轮次的代码如下：

import copy from collections import OrderedDict def fed_avg(global_model, clients, client_loaders, num_rounds=20, client_fraction=0.5): global_state = global_model.state_dict() for round_idx in range(num_rounds): # 1. 本轮参与训练的客户端采样 num_clients = max(1, int(len(clients) * client_fraction)) selected = random.sample(clients, num_clients) # 2. 每个客户端基于全局权重做本地训练 updates = [] total_samples = 0 for cid in selected: local_state = local_train( global_model, global_state, client_loaders[cid], epochs=3, lr=0.01 ) num_samples = len(client_loaders[cid].dataset) updates.append((num_samples, local_state)) total_samples += num_samples # 3. FedAvg 加权平均 new_state = OrderedDict() for key in global_state.keys(): weighted_sum = None for num_samples, local_state in updates: delta = local_state[key].float() * num_samples if weighted_sum is None: weighted_sum = delta else: weighted_sum += delta new_state[key] = (weighted_sum / total_samples).to(global_state[key].dtype) global_state = new_state global_model.load_state_dict(global_state) # 4. 每轮结束后做一次全局评估 acc = evaluate(global_model, test_loader) print(f"Round {round_idx+1}/{num_rounds}, Test Acc: {acc:.4f}") return global_state

上面weighted_sum那段就是 FedAvg 的核心。为什么要按样本数加权？因为客户端数据量有多有少，数据量大的客户端在本地训练中见过更多样本，它的权重更新应该更有话语权。

我还想强调“客户端采样”这一步。每轮只让一半客户端参与训练，不是偷懒——联邦学习面向的场景就是海量端侧设备不可能同时在线。这种随机采样机制本身也带了随机性，能起到一定的正则化作用，甚至有助于缓解 Non-IID 带来的偏差。

3.4 通信层抽象：从函数调用到 gRPC

单机模拟版本里，server 和 client 之间其实就是函数调用，可以先把逻辑打通。但如果要落地到多台机器，需要把通信层替换掉。我在项目里定义了一个最小的 gRPC 协议：

service FederatedLearning { rpc GetGlobalModel (Empty) returns (ModelWeights); rpc PushLocalUpdate (LocalUpdate) returns (Ack); } message ModelWeights { bytes weights = 1; // 序列化后的模型状态字典 } message LocalUpdate { string client_id = 1; bytes weights = 2; int32 num_samples = 3; }

实现时，客户端把state_dict()序列化成字节流上传，服务端收到后反序列化再聚合。我建议你在单机模拟跑通后再切到 gRPC，否则调试通信和训练算法交织在一起，排查问题时很难定位到底是算法问题还是网络问题。

我实际用的序列化方式是torch.save(state_dict, BytesIO())，简单可靠，不引入额外依赖。如果追求更小的传输体积，可以后续换成稀疏更新或量化编码，这个放到最后一部分展开。

4. 结果验证与对比实验：集中式、联邦 IID、联邦 Non-IID

代码写完了，接下来关键问题是：如何证明这套端云联邦方案真的有效？我设计了三个实验来回答。

4.1 实验配置与训练曲线记录

三组实验的超参数完全对齐，确保对比公平：

集中式基线整体训练 15 个 epoch；联邦实验每个客户端每轮训练 3 个 epoch，做 20 轮通信，粗略估计每个客户端总共参与过约 10 轮，累计本地 epoch 数约 30 个。这个对比其实对联邦不占优，但能真实反映端云协作在数据分散下的代价。

训练曲线我用一个 list 记录每一轮的测试准确率，最后用 matplotlib 绘制。代码很短，但视觉化之后对理解收敛过程帮助巨大：

import matplotlib.pyplot as plt def plot_curves(histories, labels): for history, label in zip(histories, labels): plt.plot(range(1, len(history)+1), history, marker='o', label=label) plt.xlabel("Communication Round") plt.ylabel("Test Accuracy (%)") plt.legend() plt.grid(True) plt.savefig("comparison.png", dpi=150)

4.2 三组实验的表现对比

我实际跑下来的结果很典型，这里说趋势不说精确数值，因为每次运行都有随机性：

• 集中式基线：准确率最高，最终收敛到约 78% 左右，收敛速度最快。
• 联邦 IID：准确率能逼近集中式，约 74% 左右，但收敛慢一些，需要十几轮通信才稳定。
• 联邦 Non-IID：准确率明显掉到 60% 上下，而且每轮之间的波动比较大。

这个结果非常好地说明了联邦学习的典型特征：数据分布越不独立同分布，FedAvg 的效果就越差。Non-IID 场景下各客户端数据分布差异太大，本地训练把模型推向不同的局部最优，服务端平均后产生冲突。波动大也是这个原因：不同轮次采样的客户端组合不同，数据偏差程度不同，聚合出的全局模型差异自然大。

如果拿这个项目去向别人展示，我最推荐展示的就是这三条曲线，因为它们直观说明了为什么联邦学习需要专门设计、为什么很多人要做 FedProx 这类改进算法。一眼就能看懂 Non-IID 是核心挑战。

4.3 隐私收益与代价的权衡

实验效果说完了，必须回到联邦学习的核心价值——隐私保护。集中式训练要把图像数据汇集到服务器，一旦服务器被攻破，数据全部暴露。端云联邦里，本地图像从未离开客户端，服务端只能看到模型权重。

但“只看权重”不等于绝对安全，这就是我强调的代价。模型权重可能反推出训练数据的统计特征，甚至是部分样本。研究者已经证明，通过梯度或权重做成员推理攻击和模型反转攻击是可能的。因此，生产环境往往要叠加差分隐私加噪或安全聚合协议，这又会对精度产生额外的损失。在你做项目规划时，一定不要忽略这个权衡。

我在实验里就做了个小验证：给权重加不同强度的高斯噪声，观察准确率下降幅度。噪声标准差 0.01 时准确率几乎无损，0.1 时直接掉 10 个百分点以上。这个数据让我对“又要隐私又要精度”的工程难度有了直观认识。

5. 常见问题排查实录与避坑清单

这部分是我真正想写的内容。前面讲的都是教科书或文档里有迹可循的东西，下面这些坑是我反复踩过才总结出来的。

5.1 训练不收敛或震荡明显

现象：全局模型的准确率一直很低（接近随机猜测），或者每轮之间上蹿下跳。

我排查的顺序依次是：

第一，看本地训练是否正常收敛。单独拿一个客户端的本地数据训练几轮，如果本地 loss 不下降，说明网络结构、学习率、数据预处理有问题，跟联邦没关系。

第二，检查聚合函数的加权是否正确。我写过一段 bug，把num_samples混进了全局验证逻辑，导致聚合权重全部错乱。建议写一个小的单元测试：构造两个客户端，上传完全相同的权重，聚合后应该还是原权重；一个客户端权重为全零，一个为全一，聚合后权重应该等于全一乘以样本占比。能用代码验证的，就别靠肉眼看。

第三，看 Non-IID 切分是否切歪了。我遇到过切分后某个客户端只有几十个样本，本地训练直接过拟合，上传的权重偏离全局模型极远。客户端样本数量必须打印出来检查。

震荡明显时，优先调低学习率。联邦场景下客户端本地学习率可以比集中式小一半，比如集中式用 0.02 时，联邦本地训练用 0.01。同时把本地训练 epoch 从 5 减到 3，减轻模型漂移。

5.2 资源占用与进程异常

CPU 环境下最容易遇到的问题是：开 10 个客户端进程后，每个客户端都默认线程数拉满，机器直接卡死。解决办法是在客户端进程入口设置torch.set_num_threads(2)或更小。

另一个常见问题是内存溢出。PyTorch CPU 版本相对省内存，但如果 DataLoader 的 num_workers 开得过大，或者多个客户端同时加载 CIFAR-10 数据集到内存，内存占用就会成倍上升。CIFAR-10 全量才一百多 MB，但每个客户端如果独立加载原始数据集再切分，10 个进程就是十几份副本。我的建议是：单机模拟时用共享内存或者让数据加载只做一次，子进程通过索引读取切分好的子集，别每个进程都重新读一遍全量数据。

5.3 数据切分带来的隐性坑

这是最容易出问题、也最不好排查的地方。

第一，类别不均衡。Non-IID 切分时如果某个客户端的样本数为 0，会在聚合时出现除零错误。我加了保护逻辑：每个客户端分配完成后检查样本数，如果低于某个阈值（比如 32），重新分配。

第二，评估偏差。全局模型下发后，我用与训练完全不同的数据做测试，这是对的。但如果测试集本身包含训练集里的样本（比如从原始训练集里切出去当测试集），准确率就虚高。CIFAR-10 官方自带测试集，直接用，别自己折腾。

第三，随机种子不一致。复现实验必须固定所有随机种子。我一开始没固定客户端采样随机种子，每次跑结果都略有不同，还以为是算法实现有误。统一设置random.seed(42)、np.random.seed(42)、torch.manual_seed(42)之后，实验才可复现。

我把最容易出现的几个问题整理成速查表，方便你对照：

6. 从模拟到落地：后续扩展思路

单机模拟跑通之后，你一定会想把方案往前推一步。我梳理了几条我认为最有价值的扩展路线。

6.1 跨机器部署与通信工程化

第一件要做的事是把 gRPC 通信层真正跑起来。Proto 定义我已经给了，下一步是启动多个客户端进程，每个进程注册到服务端。这里最大的改动是：客户端不再是“被动地被 server 调用”，而是“主动向 server 拉取模型、训练后推送更新”。实现方式上可以有两种：

• 服务端维护一个客户端连接池，等待各客户端上报，满足采样条件后广播新一轮全局模型。
• 客户端定时轮询服务端，发现模型版本更新后拉取并训练。

第二种实现更接近真实车间场景，容错也更友好。客户端断线重连后，只需要对比模型版本号，决定是否需要重新训练。

通信层需要注意心跳机制和超时重试。图像分类模型权重上传大概几 MB，在局域网环境下没问题，但跨公网时要考虑压缩和断点续传。

6.2 算法层面的改进方向

Non-IID 场景下 FedAvg 效果变差，这是几乎所有真实业务都会遇到的问题，所以我强烈建议下一步实验从这几个方向切入：

• FedProx：在本地训练损失函数中加入一项近端项，约束本地模型不要偏离全局模型太远。改动很小，只要在反向传播前给损失加上 (\frac{\mu}{2}||w - w_{global}||^2)。
• SCAFFOLD：引入控制变量来纠正客户端更新方向的偏差，实现复杂度高一些，但对极端 Non-IID 的收敛性改善更显著。
• 差分隐私：在上传权重前加高斯噪声，换取隐私保护强度。你要做的就是找出噪声尺度、隐私预算和模型精度三者之间的平衡点。
• 模型压缩：对权重做 Top-K 稀疏化或量化到 int8，减少通信量。原版联邦学习就是为带宽受限场景设计的，压缩后的效果很值得验证。

我自己测试过 FedProx，在 Non-IID 下的准确率比 FedAvg 提高了约 5 到 8 个百分点，代价是多了个超参数 (\mu) 需要调。做实验时建议先固定其它超参，只扫 (\mu) 的几个数量级（0.001、0.01、0.1），就能看出明显趋势。

6.3 工程层面的加固

算法之上，如果要往生产环境推，工程加固不能少。我的建议是加三样东西：

• 模型版本管理：每一轮全局模型都打上版本号，客户端的更新必须携带它基于的版本号。避免“旧更新覆盖新模型”的竞态问题。
• 失败重试与异常隔离：一个客户端训练到一半崩溃，不能拖垮整个训练轮次。服务端要设置超时，超时的客户端直接剔除，剩余客户端照常聚合。
• 日志与监控：给每个客户端生成结构化日志，记录本地损失、样本数、耗时。我在单机模拟阶段没有很重视日志，等到跨机器联调时才发现定位问题全靠日志。早做比晚做好。

我还想额外提一个真实体会：千万不要在模拟阶段跳过异常处理。我最初只写了“理想流程”，然后模拟客户端掉线时，聚合代码直接报错。后来补了几行 try-except 和超时判断，整个系统的稳定性立刻不一样。这些工程细节都是“用了才懂”，单机验证阶段越早补越好。

最后再分享一点个人习惯。我在跑实验时，会先把 Non-IID 的切分可视化出来——把每个客户端的数据类别分布画成热力图。这一步花不了几分钟，但能帮你一眼确认切分逻辑是否正确，避免辛辛苦苦跑完二十轮实验，最后发现数据切分一开始就错了。很多所谓“算法问题”，其实都是数据或代码层面的低级错误，可视化是最直接的排障手段。