简介:面向希望理解联邦学习完整流程的机器学习初学者与分布式训练实践者,这份Python项目围绕横向联邦学习本地模拟展开:多个虚拟客户端各自持有数据并独立训练,中心服务端负责收集参数更新、聚合全局模型并回传,帮助在没有真实多设备集群时跑通“本地训练—上传—聚合—广播”的迭代过程。压缩包共25个文件,以Python源码和XML配置为主,辅以pyc缓存、CIFAR-10数据批次及json配置等,整体约302.68MB,目录按服务端、客户端、模型、数据集与入口脚本组织,便于按模块阅读和调试。已有729人学习下载。读者可据此观察各客户端的数据分布、参数上传与平均聚合细节,理解通信接口、训练轮次控制和全局模型评估,并调整聚合策略或替换模型结构,适合作为联邦学习入门实验与课程设计的参考底稿。
1. 本地模拟横向联邦学习:先看清这份 Python 文档资料能替掉哪段工作
很多人第一次跑联邦学习模拟,会把注意力全放在聚合公式上,结果训练曲线忽上忽下,像玄学。这份用 Python 实现本地模拟横向联邦学习的文档资料,价值不在多讲一遍 FedAvg,而在把横向联邦学习拆成数据切分、客户端采样、参数聚合、评估四条可单独调试的链路;配套脚本用 Python 作为开发语言,单机就能复现客户端本地训练与服务器聚合。它适合已经写过单机训练脚本、想弄清联邦学习工程细节的人,也适合要做课程设计或方案预研的从业者。先别急着改模型,按后面的顺序把基线跑稳,再谈 Non-IID 和通信轮次。
2. 从零搭起可跑的骨架:环境、目录与启动参数
这份资源拿到手,先别翻聚合公式。先确认能不能在单机跑通。横向联邦学习的本地模拟,本质是把“服务器持有全局模型、若干客户端持有本地数据”的关系压缩到同一台机器。选型上我一般会优先用单进程对象模拟,不急着上多进程或容器编排,因为你要调的是数据切分和聚合逻辑,不是网络通信。单进程里断点、打印 state_dict、复现随机种子都直接,多进程反而把日志搅乱,出了问题连是哪个客户端返回的 loss 都难对上。
2.1 为什么本地模拟先把服务器和客户端拆成模块,而不是拆成进程
本地模拟的目标是验证算法正确性,不是压测通信开销。把服务器和客户端拆成模块,边界清楚:客户端只负责拿到全局模型副本、在本地数据上训练、返回 state_dict 和样本数;服务器只负责每轮采样客户端、按权重聚合、跑全局评估。这样你后面把 FedAvg 换成 FedProx 或加差分隐私,只动聚合模块,不用碰数据加载器。很多刚接触联邦学习的开发者喜欢把所有逻辑塞进一个train.py,前两轮能跑,等到 Non-IID 一改、客户端采样一换,变量作用域和随机种子就混在一起,曲线根本没法解释。
我常见的做法是先定四个目录:common放数据生成、切分和模型定义,client放本地训练,server放聚合和主循环,configs和logs分别放配置与运行记录。这套划分不依赖具体框架,用 Python 标准库也能落地。你拿到文档资料后,先对照目录看每个脚本的输入输出,别急着改参数;先让默认配置跑出第一轮日志,再决定要不要动客户端数量。本地模拟最怕一上来就调大 rounds,跑了半小时发现是数据切分写错了,回滚都找不到基准。
提示:先用 2 个客户端、3 轮通信跑通,确认日志里能看到每个客户端的
n_samples、local_loss和全局acc,再放大规模。
2.2 一次跑通的目录、依赖与启动命令
环境隔离先做掉,别和本机其它 Python 项目共用解释器。下面这段命令只做三件事:建虚拟环境、装最小依赖、建目录骨架。依赖里torch负责模型和 state_dict,numpy和scikit-learn负责合成数据,pyyaml读配置;如果你只用 dataclass 写配置,pyyaml可以不加。命令在 Linux、macOS 和 Windows 的 Git Bash 下都能用,纯 Windows 命令行就把mkdir -p换成逐级新建目录。
python -m venv .venv source .venv/bin/activate # Windows 用 .venv\Scripts\activate pip install torch numpy scikit-learn pyyaml mkdir -p fl_sim/{common,client,server,configs,logs}这段命令的逻辑说明很直接:虚拟环境把依赖锁在项目内,避免不同项目的torch版本互相打架;目录骨架对应“公共代码—客户端—服务器—配置—日志”五块。参数上没有需要特别调的地方,唯一要注意的是torch装 CPU 版就够本地模拟用,别为了跑几个客户端去折腾 CUDA 环境。目录建完后,启动顺序建议固定成“生成数据 → 切分客户端 → 单客户端自检 → 联邦主循环”,每一步都能单独跑,出问题不用整条链路重来。启动命令常见写法是:
python -m server.main --config configs/sim.yaml这条命令的含义是让server.main作为模块入口,读取 YAML 配置。参数--config指向配置文件,换实验时复制一份 YAML 改参数,不要直接改代码里的默认值。下面给一份最小配置,字段名你可以按资源里的实际脚本调整,但含义基本对应:
seed: 42 rounds: 30 clients: 20 clients_per_round: 5 local_epochs: 2 batch_size: 64 lr: 0.01 alpha: 0.5配置里的seed控制数据切分和客户端采样,固定它才能比较不同聚合策略;rounds是通信轮次,30 只是本地模拟的起点,不是越多越好;clients是客户端池大小,clients_per_round是每轮真正参与训练的客户端数,这两个参数别混。local_epochs控制客户端本地训练遍数,设太大客户端会往本地数据过度偏移,全局模型反而震荡;batch_size和lr要一起看,客户端样本少的时候把 batch 设到 128 以上,一个 batch 就覆盖全量本地数据,梯度噪声会很小,曲线反而失真。alpha是 Non-IID 强度,第 3 章会展开,先记住它越小客户端类别越偏。
注意:
clients_per_round不要长期等于clients。本地模拟为了快速看收敛可以全量参与,但正式对照实验要固定部分参与,否则通信轮次这个变量就失去意义。
2.3 启动参数怎么改:先固定三条基线
参数不要一次改五个。我一般先固定三条基线:第一条是单客户端本地训练,确认模型能在小数据上把 loss 降下去;第二条是两个客户端 IID 切分,确认聚合后的准确率接近单机;第三条才是 Non-IID 切分。每次只动一个参数,日志里保留round、client_id、n_samples、local_loss、global_acc五个字段。这样曲线掉下去时,你能判断是某个客户端样本太少、本地学习率太大,还是聚合权重写错。很多“联邦学习不收敛”的翻车,最后查出来是测试集和训练集重叠,或者 state_dict 的 key 顺序不一致,跟算法本身没关系。
启动参数里最容易被忽视的是seed的作用范围。只在数据生成时设numpy.random.seed不够,客户端采样如果用了另一个随机数生成器,每轮选中的客户端仍然会变。常见做法是写一个set_seed(seed),同时管random、numpy和torch,客户端采样用独立的np.random.default_rng(seed),并把每轮选中的客户端 ID 写进日志。这样复现实验时,你能确定第 7 轮选的是哪几个客户端,而不是只看一条总准确率曲线。配置快照也要跟着日志存,后面回滚有依据。
3. 横向切分与 Non-IID 数据加载器:把客户端差异做成可调参数
横向联邦的第一条边界是样本维度切分,不是特征维度切分。所有客户端共享同一套特征空间,差别在各自持有的样本不同;如果把特征列拆给不同客户端,那是纵向联邦的场景,聚合方式完全不一样。本地模拟里,我建议先用合成数据把流程跑顺,make_classification能控制类别数、特征数和噪声,复现成本低,也不会把真实数据合规问题带进实验。等你确认聚合逻辑没问题,再替换成自己的数据接口,数据加载器那一层不用重写。
3.1 横向切分的边界:切样本索引,不切特征列
下面这段代码生成一份四分类的合成数据,特征 20 维,样本 6000 条。n_informative控制真正有判别力的特征数,flip_y加一点标签噪声,避免模型一开始就 100% 准确导致看不出聚合差异。返回的X用float32,y用int64,后面直接喂给 PyTorch 的TensorDataset,不用再转类型。参数上n_samples不要设太小,客户端一多,每个客户端只剩几十条样本,本地训练一轮就过拟合;class_sep保持默认即可,本地模拟的重点不是数据难度。
import numpy as np from sklearn.datasets import make_classification def make_sim_data(n_samples=6000, n_features=20, n_classes=4, seed=42): # 合成数据只用于验证联邦流程,不涉及真实业务数据 X, y = make_classification( n_samples=n_samples, n_features=n_features, n_informative=12, n_classes=n_classes, n_clusters_per_class=1, flip_y=0.03, class_sep=1.0, random_state=seed, ) return X.astype(np.float32), y.astype(np.int64)这段代码的逻辑是先把数据难度控制在“单机模型能学到东西,但不会太完美”的区间。参数n_informative=12表示 20 维特征里有 12 维参与分类,剩下的是噪声维;flip_y=0.03表示 3% 的标签被翻转,用来观察聚合后模型对噪声的鲁棒性。random_state必须和全局 seed 对齐,否则你换一次切分就换一次数据,曲线没法横向比较。合成数据的好处是类别分布可控,坏处是别拿它的绝对准确率去推断真实场景,本地模拟看的是趋势和相对差异。
3.2 用 Dirichlet 分布构造 Non-IID 客户端
横向联邦里最常用的 Non-IID 模拟方式是按类别用 Dirichlet 分布切样本。alpha越小,每个客户端拿到的类别越集中;alpha越大,越接近 IID。下面这段函数先按类别把索引分开,每个类别打乱后按 Dirichlet 比例切给各客户端,最后合并成每个客户端的索引数组。注意alpha设到 0.1 以下时,可能出现某个客户端一条样本都没有,实际使用时要过滤空客户端,否则聚合权重为 0,state_dict 还会参与平均,容易把全局模型带偏。
def dirichlet_split(y, n_clients, alpha=0.5, seed=42): rng = np.random.default_rng(seed) client_idx = [[] for _ in range(n_clients)] for cls in np.unique(y): idx = np.where(y == cls)[0] rng.shuffle(idx) props = rng.dirichlet([alpha] * n_clients) cuts = (np.cumsum(props) * len(idx)).astype(int)[:-1] for cid, part in enumerate(np.split(idx, cuts)): client_idx[cid].extend(part.tolist()) # 过滤空客户端,避免聚合时出现零权重样本 return [np.array(sorted(x), dtype=np.int64) for x in client_idx if len(x) > 0]逻辑说明:外层循环按类别走,保证每个类别都按 Dirichlet 比例分给客户端;props是长度等于客户端数的概率向量,cumsum转成切点,np.split把当前类别的索引切成n_clients份。参数alpha=0.5是中等 Non-IID,适合第一轮对照;alpha=0.1接近极端,客户端可能只持有单一类别;alpha=100基本接近 IID,可以用来做基线。seed要和数据生成分开传,但值可以相同,方便记忆。过滤空客户端这一步别省,我见过有人在 20 个客户端、alpha=0.05下直接跑,结果 6 个客户端没有样本,聚合时权重全落在少数客户端上,曲线看着像模型崩了,其实是切分问题。
3.3 数据加载器与测试集隔离
每个客户端单独建DataLoader,不要共用一个全局 loader 再按索引过滤,后者在客户端多的时候会把内存和 IO 放大。num_workers=0是本地模拟的稳妥选择,尤其是你在 Windows 或 notebook 里跑,多进程 DataLoader 容易和模型对象、随机种子纠缠。shuffle=True让客户端本地每个 epoch 的打乱顺序不同,drop_last=False保证小客户端不会丢掉最后一批样本。测试集要在切分客户端之前就固定下来,不能先切完客户端再从训练集里抽评估数据,否则测试集泄漏会让全局准确率虚高。
import torch from torch.utils.data import TensorDataset, DataLoader def client_loader(X, y, idx, batch_size=64): ds = TensorDataset(torch.from_numpy(X[idx]), torch.from_numpy(y[idx])) return DataLoader( ds, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=False, )这段代码把索引切分和加载器解耦:X、y是全局数组,idx是当前客户端的样本索引,TensorDataset只做视图包装,不会复制整份数据。参数batch_size=64要和客户端样本量匹配;如果某个客户端只有 30 条样本,batch 设 64 会让每个 epoch 只有一个 batch,本地训练几乎变成全量梯度下降,客户端漂移会很小,但和真实联邦场景差距大。解决办法是设batch_size=min(64, len(idx)),或者在切分后过滤样本数低于阈值的客户端。测试集单独用client_loader(X_test, y_test, np.arange(len(y_test)), batch_size=256)建,shuffle=False,评估时不要打乱。
4. FedAvg 聚合与训练主循环:让每一轮通信都能被单独验证
FedAvg 的核心不是“把模型平均一下”,而是按客户端样本数加权平均。客户端数据量差异大时,简单平均会让小客户端和大客户端话语权一样,全局模型被小客户端的噪声拖偏。下面这版聚合函数遍历 state_dict 的每个 key,按weights[i] / total加权。注意它默认所有 state_dict 的 key 顺序一致、张量都是浮点;如果模型里有 BatchNorm 的num_batches_tracked这种整型 buffer,直接加权会变成浮点,load_state_dict时报类型错误。第 5 章会把这个问题展开。
4.1 聚合公式与加权方式
import copy def fedavg(states, weights): total = sum(weights) avg = copy.deepcopy(states[0]) for k in avg.keys(): avg[k] = sum( states[i][k] * (weights[i] / total) for i in range(len(states)) ) return avg逻辑说明:states是每轮参与训练的客户端返回的 state_dict 列表,weights是对应客户端的样本数。先深拷贝第一个状态作为模板,再对每个 key 做加权求和。参数weights不要传成每个客户端的 loss 或准确率,FedAvg 的标准权重是样本数;如果你要换加权策略,建议在函数外面算好权重再传进来,聚合函数保持单一职责。张量设备也要统一,本地模拟一般在 CPU 上聚合,再把聚合后的 state_dictload_state_dict到全局模型;如果全局模型在 GPU 上,先load_state_dict再.to(device),别在聚合时混着 GPU 和 CPU 张量做加法。
4.2 客户端本地训练函数
客户端本地训练要拿全局模型的副本,不能直接改全局模型。copy.deepcopy(model)每轮每个客户端一次,客户端多时内存会涨,但本地模拟规模不大,先保证正确性。优化器用 SGD 还是 Adam 会影响客户端漂移程度,FedAvg 原始实验常用 SGD,本地模拟想复现论文趋势就用 SGD;如果只关心能不能跑通,Adam 收敛快,但别拿 Adam 的结果去对比别人的 SGD 基线。返回值里必须带样本数,服务器聚合时要用。
import torch.nn as nn def local_train(model, loader, epochs, lr, device): local_model = copy.deepcopy(model).to(device) opt = torch.optim.SGD(local_model.parameters(), lr=lr) loss_fn = nn.CrossEntropyLoss() local_model.train() for _ in range(epochs): for xb, yb in loader: xb, yb = xb.to(device), yb.to(device) opt.zero_grad() loss = loss_fn(local_model(xb), yb) loss.backward() opt.step() return local_model.state_dict(), len(loader.dataset)逻辑说明:每个客户端从同一份全局模型出发,在本地数据上跑epochs遍,返回更新后的 state_dict 和本地样本数。参数epochs对应配置里的local_epochs,设 1 到 2 先看趋势;lr是客户端本地学习率,不是服务器学习率,FedAvg 里服务器通常没有额外学习率,聚合完直接替换全局模型。device在本地模拟里用 CPU 就够,xb、yb每次 batch 都.to(device),别在 DataLoader 里提前搬到 GPU,客户端多的时候显存会被重复占用。
4.3 服务器主循环与评估
主循环负责三件事:每轮采样客户端、收集本地状态、聚合后评估。采样用replace=False,保证同一轮不重复选同一个客户端;clients_per_round小于clients时才像联邦通信,等于全量参与则更像分布式 SGD。评估函数要model.eval()加torch.no_grad(),否则 BatchNorm 和 Dropout 会处在训练模式,准确率每次都在抖。日志至少打印轮次、参与客户端 ID、聚合后的全局准确率,调试阶段再加每层权重的 L2 范数,看聚合有没有把参数放大。
def evaluate(model, loader, device): model.eval() correct = total = 0 with torch.no_grad(): for xb, yb in loader: xb, yb = xb.to(device), yb.to(device) pred = model(xb).argmax(dim=1) correct += (pred == yb).sum().item() total += yb.numel() return correct / max(total, 1)主循环的伪代码可以写成下面这样,实际脚本里把数据、模型、优化器初始化补全即可。参数rounds控制通信轮次,clients_per_round控制每轮参与数,local_epochs和lr从配置读。每轮聚合后立刻评估,别等所有轮跑完再补评估,否则中间哪轮开始发散你根本看不到。如果第 1 轮准确率就是随机水平,先看客户端本地 loss 有没有下降,再看聚合权重是不是全给了空客户端。
for r in range(rounds): selected = rng.choice(n_clients, clients_per_round, replace=False) states, weights = [], [] for cid in selected: sd, n = local_train(global_model, loaders[cid], local_epochs, lr, device) states.append(sd) weights.append(n) new_sd = fedavg(states, weights) global_model.load_state_dict(new_sd) acc = evaluate(global_model, test_loader, device) print(f"round={r} clients={selected.tolist()} acc={acc:.4f}")5. 避坑与排查:五类高频翻车现场和定位顺序
本地模拟联邦学习最耗时的不是写模型,是判断曲线不对时该先看哪。下面五条是血泪经验里出现频率最高的,按“现象→原因→解决”写。排查时不要一次改多个参数,先把数据切分和单客户端基线钉死,再动聚合逻辑。
5.1 五类高频翻车现场
第一类:聚合后准确率比单客户端基线还低。现象是第 1 轮全局准确率还有 0.7,第 3 轮掉到 0.28,看起来像聚合把模型带崩了。原因通常有三个:客户端样本量差异大,简单平均让只有几十条样本的客户端和大客户端同权;各客户端返回的 state_dict key 顺序不一致,聚合时按位置相加,实际加错了层;测试集和某个客户端的训练集重叠,单客户端基线虚高。解决方法是按n_samples加权,聚合前断言所有 state_dict 的 key 列表相同,测试集在切分客户端之前单独留出。先跑两个客户端、IID 切分,确认聚合后不比单客户端差,再换 Non-IID。
第二类:local_loss 变成 nan 或梯度爆炸。现象是某个客户端训练到第二个 epoch 时 loss 直接 inf,聚合后全局模型输出全是 nan。原因多半是本地学习率相对 batch 太大、合成数据特征没标准化、或者某个客户端样本太少导致梯度方向极端。解决方法是先把X做标准化,lr从 0.01 降到 0.001 试一轮,batch_size不要小于 32;再加梯度裁剪torch.nn.utils.clip_grad_norm_(local_model.parameters(), 5.0)。本地模拟里先单独跑这个客户端,看它自己能不能把 loss 降下去,别直接怀疑聚合公式。
第三类:同一配置每次跑出来的准确率差十几个点。现象是配置没改,第一次 0.82,第二次 0.69。原因通常是 seed 只固定了数据生成,没固定客户端采样和 DataLoader 的 shuffle;或者你用了全局随机数生成器,数据切分和客户端采样互相影响。解决方法是写统一的set_seed,把random、numpy、torch都管住,客户端采样用独立的default_rng(seed),每轮把选中的客户端 ID 写进日志。配置文件和 git commit 一起存,回滚时才能对上。
第四类:DataLoader 开多进程后卡死或内存持续上涨。现象是num_workers调到 2 以上,程序跑几轮就挂起,或者内存一条直线往上走。原因是本地模拟里模型和数据集对象被 fork 复制,每个客户端 loader 又各开一组进程,客户端一多进程数爆炸。解决方法是本地模拟统一num_workers=0,客户端串行训练;数据张量只保留一份,用索引取 batch,不要给每个客户端复制整份X。如果确实要并行,先限制客户端数量,再逐步加 worker,别一上来就拉满。
第五类:全局准确率虚高,客户端本地评估却很差。现象是全局测试集 0.99,但把全局模型发回每个客户端本地评估,准确率只有 0.6。原因通常是测试集泄漏:切分客户端时把测试集索引也切进去了,或者用测试集调过alpha和rounds。解决方法是先留出固定测试集,再在剩余数据上做客户端切分;测试集只用于最终评估,调参用验证集。再补一个混淆矩阵,看全局准确率高是不是因为某一类样本特别多,把多数类预测对了。
5.2 排查顺序:先数据,再单客户端,再聚合
曲线不对时,我一般按这个顺序走:第一步打印每个客户端的样本量和类别分布,确认没有空客户端、没有某个客户端独占某一类;第二步拿一个客户端单独训练,看它能不能在本地数据上过拟合, loss 不降就先查学习率和特征标准化;第三步用两个客户端做 IID 切分,聚合后应该接近单机;第四步换 Non-IID,观察掉点幅度;最后才调rounds和模型结构。这个顺序能避免你在聚合公式上反复改,最后发现是数据切分写错了。
for cid, idx in enumerate(client_idx): labels, counts = np.unique(y[idx], return_counts=True) print(cid, "n=", len(idx), dict(zip(labels.tolist(), counts.tolist())))这段检查代码逐个客户端打印样本数和类别计数,参数client_idx是第 3 章切分函数的返回值。逻辑上它不参与训练,只做数据体检;如果某个客户端n小于batch_size,要么调小 batch,要么在
本文还有配套的精品资源,点击获取