☰
单机横向联邦学习实战:FedAvg聚合与Non-IID切分
2026/11/2 2:07:05 网站建设 项目流程

简介:面向希望理解横向联邦学习完整流程的Python开发者与机器学习学习者,这套资料以本地模拟方式搭建多客户端协同训练环境,重点解决不共享原始数据前提下的模型训练、参数上传与服务器聚合问题,适合具备一定深度学习基础、想动手复现联邦学习通信与聚合机制的人群。压缩包共25个文件,约302.68MB,以Python源码和编译缓存为主,辅以XML/IDE配置、JSON配置、CIFAR-10数据批次及压缩数据文件,覆盖客户端训练、服务端协调、模型定义、数据集加载和运行入口等环节。已有729人学习下载,说明该主题在隐私保护与分布式机器学习实践中具有一定关注度。读者可据此梳理本地模拟联邦学习的工程结构,观察多客户端训练、参数传输、全局聚合与迭代更新的实现脉络,并借助自带数据与配置快速跑通实验,为后续替换模型、调整聚合策略或加入加密与差分隐私等扩展提供可复用的代码框架。

1. 单机跑通横向联邦:为什么把 N 个客户端塞进一个进程反而更接近真相

一台笔记本、一份 Python 环境,想验证横向联邦学习里 FedAvg 到底怎么收敛,真去搭多机服务是最笨的路。这份资源把 N 个客户端塞进同一个进程里跑,用函数调用代替网络通信,重点全压在非独立同分布切分、客户端采样和加权聚合这三件事上。它适合正在做课程设计、要把联邦学习写进毕设,或者面试前想把“聚合权重到底按谁的样本数算”讲清楚的人。开发语言只用 Python,依赖很薄,跑通一轮不到一分钟,改一个参数就能看到全局模型震荡——这种即时反馈在真机集群上很难拿到。

2. 拆开 FedAvg 的骨架:从样本维度切分到二十行加权聚合

横向联邦这个词被用得很泛,先把边界划清楚再谈代码。横向联邦(horizontal federated learning)指的是各参与方特征空间一致、样本 ID 不重叠,切的是样本维度;对应的纵向联邦切的是特征维度,需要样本对齐,实现复杂度完全不是一个量级。本地模拟要复现的是横向那一支,也就是“同样的模型结构,各自喂自己的数据,只交换参数”。

2.1 横向联邦切的是样本,不是特征

假设有一份带标签的数据集,10 万条样本、10 个类别、特征维度固定。横向联邦的做法是把这 10 万条按某种规则分给 K 个客户端,每个客户端拿到的仍然是完整的特征列,只是行数不同、类别分布可能不同。全局模型始终只有一份结构,客户端本地训练出来的也只是同一份结构的另一组参数。

这个定义决定了三件事:模型结构必须全客户端一致,否则参数没法逐元素相加;聚合只能发生在参数层,不能发生在梯度层(梯度聚合等于把本地 epoch 压到 1,那是另一种算法);评估必须有一份谁都没见过的测试集,否则你测的是聚合后的过拟合程度。

纵向联邦那套样本对齐、同态加密、隐私求交的东西在这里完全用不上。看到有人在横向联邦的模拟代码里塞 PSI 协议,基本可以判断是把两个场景抄混了。

2.2 单机模拟的进程模型:函数调用代替网络往返

真机上客户端是独立进程甚至独立机器,一轮通信要序列化参数、走网络、反序列化。本地模拟把这层砍掉,换成一次deepcopy加一次函数返回。这个替换带来两个好处和一个必须警惕的副作用。

好处一是调试成本骤降,客户端本地训练报错直接就是当前进程的栈,不用去翻容器日志;好处二是可以在一分钟内跑完 100 轮,参数扫描做起来才有意义。副作用是对象引用会被共享——真机上参数过网络一定是副本,本地模拟如果不显式深拷贝,客户端拿到的就是全局模型参数的同一个引用,本地训练会直接改掉全局模型,“聚合”变成一句空话。这一点在后面的避坑章节还会展开。

工程目录一般是这个形状,文件不多,但每个都得有明确职责:

文件职责常见误用
partition.py按 Dirichlet 或 IID 切分索引,落盘成 json每次训练重新切,导致实验不可复现
client.py封装本地 Dataset、DataLoader 与 local_train把 optimizer 建在全局,状态跨轮残留
server.py客户端采样、加权聚合、全局评估聚合权重用全局样本数而不是本轮样本数
train.py主循环、随机种子、指标落盘种子只固定了 torch,漏了 numpy 和 random
configs/*.yaml轮次、本地 epoch、α、学习率参数写死在代码里,扫描实验时改到崩溃

2.3 二十行 numpy 实现 FedAvg 的加权聚合

先把聚合函数单独抽出来跑通,再谈训练。FedAvg 的更新式是 $w_{t+1} = \sum_k \frac{n_k}{n} w_t^k$,其中 $n_k$ 是客户端 k 本轮参与训练的样本数,$n$ 是本轮所有参与客户端的样本总数。写成代码就二十行,但权重口径错了整个实验就废了。

import numpy as np def fedavg_aggregate(client_states, client_sizes): """FedAvg 加权聚合。 client_states: List[Dict[str, np.ndarray]],每个客户端本轮训练后的参数字典 client_sizes: List[int],与 client_states 一一对应的本轮样本数 n_k 返回: Dict[str, np.ndarray],新的全局参数 """ if len(client_states) != len(client_sizes): raise ValueError("参数字典与样本数列表长度不一致") total = int(sum(client_sizes)) if total == 0: # 采样比例过小或某个客户端分到 0 条样本时会踩到 raise ValueError("本轮参与样本总数为 0,检查客户端采样与数据划分") new_state = {} for key in client_states[0].keys(): acc = np.zeros_like(client_states[0][key], dtype=np.float64) for state, n in zip(client_states, client_sizes): acc += state[key].astype(np.float64) * n # 按样本数加权,不是简单平均 new_state[key] = (acc / total).astype(client_states[0][key].dtype) return new_state

逻辑上有三个点值得说清楚。第一,累加用float64中转再转回原 dtype,是因为 float16 或 bfloat16 在累加几十个客户端的参数时会丢精度,全局模型会出现肉眼可见的漂移,这种问题排查起来极其费劲。第二,分母total是本轮参与客户端的样本数之和,不是全体客户端的样本总数——如果按全体算,参与率 0.3 的时候聚合结果会被系统性缩小,等价于给学习率乘了个 0.3。第三,键的集合取自第一个客户端,前提是所有客户端模型结构一致,这一点在横向联邦里是硬约束,可以在聚合前加一个assert set(s.keys()) == set(client_states[0].keys())做防守。

PyTorch 版本只需要把 numpy 换成 tensor,其余逻辑照抄:

import torch def fedavg_torch(client_states, client_sizes): total = float(sum(client_sizes)) new_state = {} for key in client_states[0]: acc = torch.zeros_like(client_states[0][key], dtype=torch.float32) for state, n in zip(client_states, client_sizes): acc += state[key].float() * (n / total) # 先归一化再乘,数值更稳 new_state[key] = acc.to(client_states[0][key].dtype) return new_state

2.4 参数先定死:轮次、本地 epoch 与学习率的取值区间

新手最容易犯的错是一上手就把轮次开到 500、本地 epoch 开到 10,然后盯着不收敛的曲线怀疑算法。先给一组能出结果的保守值,跑通再放大:

参数建议初值取值范围调大的后果
rounds5020 ~ 200收益递减,后 100 轮往往只涨零点几个点
local_epochs11 ~ 5超过 5 客户端漂移明显,全局曲线开始锯齿
batch_size3216 ~ 128太大导致单轮本地更新步数太少,等于没训
lr0.01(SGD)0.001 ~ 0.050.05 以上在 Non-IID 下极易发散
frac(参与率)0.50.1 ~ 1.0太低梯度噪声大,曲线抖到看不出趋势

本地 epoch 这个参数值得单独说。它代表“客户端在每轮里把本地数据过几遍”,过多了客户端参数会跑向各自的最优解,聚合回全局时互相拉扯,这就是常说的客户端漂移(client drift)。曲线上的典型表现是准确率上去一点又掉回来,来回震荡,看着像玄学,其实是local_epochs和lr都太大了。

提示:先用local_epochs=1、frac=1.0、IID 划分跑一遍,确认能收敛到接近集中式训练的水平,再逐个把参数改成你要研究的设定。基线不稳,后面所有结论都是假的。

3. 非独立同分布切分:Dirichlet 的 α 到底该取多少

联邦学习论文里绝大多数结论都建立在“数据是 Non-IID 的”这个前提上。如果模拟器把数据均匀打散分给客户端,跑出来的 FedAvg 曲线会漂亮得不像话,然后你拿着这个结果去对比 FedProx,发现两者几乎一样——因为场景本身没有难度。

3.1 IID 切分和 Non-IID 切分差在哪

IID 划分离谱地简单:全局打乱索引,按 K 等分。每个客户端拿到的类别比例和全局一致,本地训练的梯度方向也高度一致,聚合近似于把所有数据集中起来训。这种设定下 FedAvg 和集中式训练的差距通常在 1 个点以内,用来做正确性自检很合适,用来做算法对比毫无意义。

Non-IID 的现实来源是各参与方的业务不同。比如同一个手写数字识别任务,有的客户端扫的都是某几个数字,有的客户端只覆盖另外几个类别,而且样本量还严重不均衡。这类分布在代码里通常用 Dirichlet 分布来生成:对每个类别,从一个参数为 α 的 Dirichlet 分布里采样 K 个比例,按比例把该类样本切给各客户端。

α 越小,比例越极端,越接近“每个客户端只拿到一两个类别”(病态 Non-IID);α 越大,比例越均匀,α=100 时基本退化成 IID。工程上常用 0.1、0.5、1.0 三档做对比,其中 0.1 是最能拉开算法差距的档位,也最容易把没写好的聚合逻辑打回原形。

3.2 Dirichlet 切分脚本与保底名额

下面这个划分函数里有个容易被忽略的细节:min_size。当 α=0.1 时,Dirichlet 采样会给出接近 0 的比例,某个客户端在某个类别上一条样本都分不到是很常见的;如果所有类别都这样,就会出现空客户端,后面的样本数加权直接除零。

import numpy as np def dirichlet_split(labels, num_clients, alpha=0.5, min_size=10, seed=0): """按类别做 Dirichlet 切分,返回每个客户端的样本索引列表。 labels: 一维数组,全局数据集每条样本的类别 num_clients: 客户端数量 K alpha: 越小越异构,0.1 极异构 / 0.5 中等 / 100 近似 IID min_size: 每个客户端在每个类别上的保底样本数,防止空客户端 """ rng = np.random.default_rng(seed) labels = np.asarray(labels) num_classes = int(labels.max()) + 1 client_idx = [[] for _ in range(num_clients)] for c in range(num_classes): idx = np.where(labels == c)[0] rng.shuffle(idx) # 先打乱,否则切出来的是有序片段 props = rng.dirichlet(np.repeat(alpha, num_clients)) budget = len(idx) - min_size * num_clients if budget >= 0: sizes = np.floor(props * budget).astype(int) # floor 会丢余数,补到配额最大的那个客户端上,保证总数对得上 sizes[np.argmax(sizes)] += budget - sizes.sum() sizes = sizes + min_size # 再把保底名额加回来 else: # 该类样本太少,退化成尽量均分 sizes = np.full(num_clients, len(idx) // num_clients) sizes[: len(idx) % num_clients] += 1 start = 0 for k, s in enumerate(sizes): client_idx[k].extend(idx[start:start + s].tolist()) start += s return client_idx

参数上,seed必须固定并且和训练用的种子分开管理,因为划分结果通常要落盘复用,两次实验如果划分不同,指标差异根本没法归因。min_size取 10 是个经验值,样本量小的数据集(比如每类只有几百条)要相应调到 3 到 5,否则budget直接变负,代码走退化分支,切出来的分布反而接近 IID,你会误以为 α 没起作用。

切完把索引存成 json,训练时只读不重算:

import json client_idx = dirichlet_split(labels, num_clients=10, alpha=0.1, seed=20240101) with open("partitions/dir0.1_seed20240101.json", "w") as f: json.dump(client_idx, f)

3.3 客户端采样:每轮选谁,权重按什么算

真实的横向联邦里,客户端不会每轮都在线。模拟器用frac控制参与率,每轮从 K 个客户端里无放回抽m = max(int(frac * K), 1)个。这里有个容易被忽略的点:抽到的客户端集合每轮都在变,聚合分母必须跟着变,用本轮参与者的样本总数。

def sample_clients(num_clients, frac, rng): m = max(int(frac * num_clients), 1) return rng.choice(num_clients, size=m, replace=False).tolist()

另一个常见做法是把“每轮随机抽”换成“固定子集轮转”,也就是把所有客户端切成若干组,按组轮换。这种做法在写论文对比实验时更受青睐,因为消除了采样带来的方差,两个算法之间的差异更容易归因到算法本身而不是这一轮恰好抽到了谁。

3.4 划分完先做一次体检

切完数据别急着开训,先打印每个客户端的样本数和类别分布。健康的 α=0.1 划分应该看到:样本数从几十到上万不等,类别分布明显倾斜,有的客户端只有 2 到 3 个类别。如果发现每个客户端样本数几乎相等、类别比例也差不多,那基本是 α 太大或者min_size撑爆了 budget,切分退化成了均分。

体检代码很简单,几行就够:

for k, idx in enumerate(client_idx): cls, cnt = np.unique(labels[idx], return_counts=True) print(f"client {k:02d} | n={len(idx):5d} | classes={dict(zip(cls.tolist(), cnt.tolist()))}")

看到某个客户端n=0,直接回头查min_size和类别样本量,不要往下跑。空客户端在聚合时贡献 0 权重,表面不报错,实际等于这一轮的参与率比你以为的低,实验结论全偏。

4. 训练主循环:state_dict 搬运、本地 epoch 与指标落盘

骨架和数据都齐了,剩下的就是把广播、本地训练、聚合、评估串成一个循环。这一步的代码量不大,但每一行都有对应的坑,写错一处曲线就废。

4.1 参数搬运的核心是深拷贝

PyTorch 的state_dict()返回的是张量的引用,不是副本。如果直接把global_model.state_dict()交给客户端,客户端load_state_dict之后做本地训练,优化器更新的就是全局模型的那批参数对象,服务端手里的“全局模型”已经悄悄变成了最后一个客户端的本地模型。

正确做法有两种,选一种并坚持:一是每次从全局模型copy.deepcopy出一份新模型给客户端;二是客户端训练完把参数detach().cpu().clone()出来再回传。两种都行,混用也不会出错,但绝对不能不拷贝。

import copy def snapshot(model): """取一份与全局模型完全脱钩的参数快照,CPU 上的克隆。""" return {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} local_model = copy.deepcopy(global_model) # 结构和参数都是新的

4.2 本地训练函数与主循环

本地训练函数只做一件事:在给定的 DataLoader 上跑固定轮数,返回参数。优化器必须在函数内部新建,不能提到外面复用,否则动量缓冲会跨轮、跨客户端累积,等于给每个客户端偷偷接了上一个客户端的训练状态。

import torch import torch.nn as nn def local_train(model, loader, epochs, lr, device): model.train() model.to(device) opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) crit = nn.CrossEntropyLoss() for _ in range(epochs): for x, y in loader: x, y = x.to(device), y.to(device) opt.zero_grad() loss = crit(model(x), y) loss.backward() opt.step() return snapshot(model)

主循环按“采样 → 本地训练 → 加权聚合 → 周期评估”的顺序推进,评估频率用eval_every控制,没必要每轮都跑测试集,测试集大的时候这一项能占掉一半时间。

rng = np.random.default_rng(seed) global_model = build_model().to(device) for rnd in range(1, rounds + 1): participants = sample_clients(num_clients, frac, rng) states, sizes = [], [] for cid in participants: local_model = copy.deepcopy(global_model) state = local_train(local_model, client_loaders[cid], local_epochs, lr, device) states.append(state) sizes.append(len(client_datasets[cid])) # 本轮该客户端的样本数 global_model.load_state_dict(fedavg_torch(states, sizes)) if rnd % eval_every == 0 or rnd == rounds: acc, loss = evaluate(global_model, test_loader, device) writer.write({"round": rnd, "acc": acc, "loss": loss, "participants": len(participants), "samples": sum(sizes)})

load_state_dict默认strict=True,键对不上会直接抛错,这是好事,别为了“跑通”改成strict=False——那等于把结构不一致的问题藏起来,等你发现某些层根本没被更新时已经跑了几十轮。

4.3 指标落盘与曲线怎么读

指标别只 print,落成 jsonl 一行一轮,方便后面画收敛曲线和算通信轮次到目标精度(rounds-to-accuracy)。这个指标比最终精度更能反映算法效率,在联邦场景里通信轮次就是成本。

曲线读法上有几个固定套路:IID 设定下曲线应该平滑单调上升,出现锯齿说明学习率大了或者 batch 太小;Non-IID 下前 20 轮掉点严重是正常的,因为初始几轮的客户端分布和测试集差得远;如果准确率在某一轮突然跳到 0.1(10 分类的随机水平),多半是聚合后的参数被某个客户端的 NaN 污染了,去查那个客户端的学习率和数据里有没有脏样本。

注意:evaluate一定要用model.eval()加torch.no_grad(),并且评估前后不要改 BatchNorm 的 running stats。评估写错导致的“精度虚高”是最难发现的一类 bug,因为曲线看起来很健康。

5. 避坑排查:五条把模型训崩的血泪记录

这一章是我在复现过程中真实踩过的五条,按“现象 → 原因 → 解决”记下来,遇到对应的症状可以直接对照。

5.1 数据与划分层的三个坑

坑一:每轮准确率纹丝不动,像没训练一样。现象是 50 轮跑完,准确率一直在初始水平附近小幅波动。原因大概率是client_datasets[cid]每次返回的都是同一个对象,或者 DataLoader 的shuffle=False加上数据本身按类别排序,导致每个 batch 全是同一类。解决办法是把 Dataset 建成只读,DataLoader 开shuffle=True,并给generator单独传一个torch.Generator().manual_seed(seed + cid),否则多个 DataLoader 会共用全局随机状态,采样顺序互相干扰。

坑二:α 设成 0.1,结果和 α=1.0 差不多。现象是无论怎么调 α,各客户端的类别分布都很均匀。原因是min_size设得太大:每类样本只有 600 条、客户端 10 个、min_size=10时保底名额就吃掉 100 条,剩下的预算只够轻微扰动,Dirichlet 的极端性被抹平了。解决是把min_size压到 3 到 5,或者干脆设成 0 再单独过滤空客户端。判断依据很简单,打印每个客户端的类别直方图,看最大类别占比是否超过 60%。

坑三:换个随机种子,结论反转。现象是同一套超参,种子 42 比种子 2024 高出 8 个点。原因是只固定了torch.manual_seed,漏了np.random.seed和random.seed,更要命的是数据划分用的default_rng(seed)和训练用的种子不是同一个来源,导致“划分固定了但采样没固定”。解决是把划分种子和训练种子写进 config 分别管理,划分结果落盘复用,训练侧在入口处统一设三个种子:

import random, numpy as np, torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False

5.2 训练与聚合层的两个坑

坑四:全局模型评估崩到随机水平,训练 loss 却正常。现象是本地 loss 一路降,聚合后测试准确率只有 10%。原因是模型里有 BatchNorm,而running_mean、running_var这些 buffer 在 Knife 边缘上根本没被聚合——我的fedavg_torch只遍历了state_dict(),理论上 buffer 也在里面,但客户端的 buffer 只由本地那点 Non-IID 数据统计出来,聚合后和全局权重严重错配。解决有两条路:把 BN 换成 GroupNorm 或 LayerNorm,彻底摆脱对 batch 统计的依赖;或者在聚合时对 buffer 用简单平均而不是样本数加权,并在评估前用一小部分无标签数据重新校准。我一般选前者,省事。

坑五:曲线前 30 轮很好,后面开始持续下滑。现象是准确率冲到峰值后缓慢掉头向下,像过拟合但训练集准确率也没涨。原因是local_epochs设成了 5 且学习率 0.05,客户端在本地把参数拉得太远,聚合时互相抵消,等效于往全局模型里注入噪声。解决是把local_epochs降到 1 到 2,或者保持 epoch 但加近端项约束本地参数不要离全局太远。判断依据是看每轮聚合前后参数的 L2 距离,距离持续增大就说明漂移在累积。

6. 从能跑到能比:给模拟器加近端项与三步自检

一个只能跑 FedAvg 的模拟器价值有限,真正有用的是它能让你在同一份划分、同一个种子、同一套评估协议下对比不同算法。做到这一步只需要加一个近端项和一套自检流程。

6.1 加十行代码把 FedAvg 变成 FedProx

FedProx 的思路是在本地损失里加一项 $\frac{\mu}{2}|w - w^t|^2$,惩罚本地参数偏离本轮全局参数的程度。改动只发生在local_train里,主循环一行不用动。

def local_train_prox(model, loader, epochs, lr, device, global_params, mu=0.01): model.train() opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) crit = nn.CrossEntropyLoss() for _ in range(epochs): for x, y in loader: x, y = x.to(device), y.to(device) opt.zero_grad() loss = crit(model(x), y) prox = 0.0 for p, p0 in zip(model.parameters(), global_params): prox = prox + ((p - p0) ** 2).sum() loss = loss + mu / 2.0 * prox # mu 越大越贴近全局,0 等价于 FedAvg loss.backward() opt.step() return snapshot(model)

global_params必须是本轮开始前从全局模型克隆出来的一批detach()张量,不能传模型对象,否则反向传播会顺着近端项把梯度也灌回全局模型。mu的常用档位是 0.001、0.01、0.1,α 越小时mu该越大,但超过 0.1 之后本地基本训不动,曲线会变得和只做聚合不做本地训练差不多。这一段写完,configs/里加一个algo: fedavg | fedprox的开关,对比实验台就成了。

6.2 三步自检:确认你的模拟器没写错

在拿它跑任何对比之前,按这三步走一遍,任何一步不过都不要继续。

第一步,单客户端退化测试。把num_clients设成 1、frac设成 1.0,此时联邦训练应该和集中式训练逐轮几乎完全一致。如果对不上,问题在本地训练函数或评估逻辑,不在聚合。

第二步,IID 上限测试。alpha=100、frac=1.0、local_epochs=1,跑 50 轮,最终精度和集中式训练的差距应该在 1 到 2 个点以内。差距过大说明聚合权重口径错了,回头查分母。

第三步,Non-IID 压力测试。alpha=0.1、frac=0.3,此时曲线应该明显比 IID 抖动,最终精度低 5 到 15 个点。如果几乎没差别,说明数据划分没生效,回去看 3.4 的体检输出。

测试配置期望现象不通过先查
退化K=1, frac=1.0与集中式逐轮一致本地训练、评估
IID 上限α=100, frac=1.0差距 ≤ 2 点聚合分母
Non-IID 压力α=0.1, frac=0.3抖动明显、低 5~15 点划分是否生效

三步都过,再把algo切到fedprox,用同一份划分文件、同一个种子跑一遍,两条曲线放在一张图里。这时候你看到的差距才是算法本身带来的,而不是随机种子或划分差异。这套流程我现在每接一个联邦学习相关的活儿都会强制走一遍,因为早年间有一次没做退化测试,聚合分母写成了全局样本数,跑了两周才发现所有对比结论都是参与率的函数。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询