资讯详情

Python本地模拟横向联邦学习:FedAvg聚合与Non-IID切分实战

📅 2026/11/2 15:34:47 | 华诺云谱 👁 阅读
Python本地模拟横向联邦学习:FedAvg聚合与Non-IID切分实战
简介面向希望理解联邦学习完整流程的机器学习初学者与分布式训练实践者这份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_dictnumpy和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用float32y用int64后面直接喂给 PyTorch 的TensorDataset不用再转类型。参数上n_samples不要设太小客户端一多每个客户端只剩几十条样本本地训练一轮就过拟合class_sep保持默认即可本地模拟的重点不是数据难度。import numpy as np from sklearn.datasets import make_classification def make_sim_data(n_samples6000, n_features20, n_classes4, seed42): # 合成数据只用于验证联邦流程不涉及真实业务数据 X, y make_classification( n_samplesn_samples, n_featuresn_features, n_informative12, n_classesn_classes, n_clusters_per_class1, flip_y0.03, class_sep1.0, random_stateseed, ) return X.astype(np.float32), y.astype(np.int64)这段代码的逻辑是先把数据难度控制在“单机模型能学到东西但不会太完美”的区间。参数n_informative12表示 20 维特征里有 12 维参与分类剩下的是噪声维flip_y0.03表示 3% 的标签被翻转用来观察聚合后模型对噪声的鲁棒性。random_state必须和全局 seed 对齐否则你换一次切分就换一次数据曲线没法横向比较。合成数据的好处是类别分布可控坏处是别拿它的绝对准确率去推断真实场景本地模拟看的是趋势和相对差异。3.2 用 Dirichlet 分布构造 Non-IID 客户端横向联邦里最常用的 Non-IID 模拟方式是按类别用 Dirichlet 分布切样本。alpha越小每个客户端拿到的类别越集中alpha越大越接近 IID。下面这段函数先按类别把索引分开每个类别打乱后按 Dirichlet 比例切给各客户端最后合并成每个客户端的索引数组。注意alpha设到 0.1 以下时可能出现某个客户端一条样本都没有实际使用时要过滤空客户端否则聚合权重为 0state_dict 还会参与平均容易把全局模型带偏。def dirichlet_split(y, n_clients, alpha0.5, seed42): 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), dtypenp.int64) for x in client_idx if len(x) 0]逻辑说明外层循环按类别走保证每个类别都按 Dirichlet 比例分给客户端props是长度等于客户端数的概率向量cumsum转成切点np.split把当前类别的索引切成n_clients份。参数alpha0.5是中等 Non-IID适合第一轮对照alpha0.1接近极端客户端可能只持有单一类别alpha100基本接近 IID可以用来做基线。seed要和数据生成分开传但值可以相同方便记忆。过滤空客户端这一步别省我见过有人在 20 个客户端、alpha0.05下直接跑结果 6 个客户端没有样本聚合时权重全落在少数客户端上曲线看着像模型崩了其实是切分问题。3.3 数据加载器与测试集隔离每个客户端单独建DataLoader不要共用一个全局 loader 再按索引过滤后者在客户端多的时候会把内存和 IO 放大。num_workers0是本地模拟的稳妥选择尤其是你在 Windows 或 notebook 里跑多进程 DataLoader 容易和模型对象、随机种子纠缠。shuffleTrue让客户端本地每个 epoch 的打乱顺序不同drop_lastFalse保证小客户端不会丢掉最后一批样本。测试集要在切分客户端之前就固定下来不能先切完客户端再从训练集里抽评估数据否则测试集泄漏会让全局准确率虚高。import torch from torch.utils.data import TensorDataset, DataLoader def client_loader(X, y, idx, batch_size64): ds TensorDataset(torch.from_numpy(X[idx]), torch.from_numpy(y[idx])) return DataLoader( ds, batch_sizebatch_size, shuffleTrue, num_workers0, drop_lastFalse, )这段代码把索引切分和加载器解耦X、y是全局数组idx是当前客户端的样本索引TensorDataset只做视图包装不会复制整份数据。参数batch_size64要和客户端样本量匹配如果某个客户端只有 30 条样本batch 设 64 会让每个 epoch 只有一个 batch本地训练几乎变成全量梯度下降客户端漂移会很小但和真实联邦场景差距大。解决办法是设batch_sizemin(64, len(idx))或者在切分后过滤样本数低于阈值的客户端。测试集单独用client_loader(X_test, y_test, np.arange(len(y_test)), batch_size256)建shuffleFalse评估时不要打乱。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(), lrlr) 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 服务器主循环与评估主循环负责三件事每轮采样客户端、收集本地状态、聚合后评估。采样用replaceFalse保证同一轮不重复选同一个客户端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(dim1) 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, replaceFalse) 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(fround{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_workers0客户端串行训练数据张量只保留一份用索引取 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_countsTrue) print(cid, n, len(idx), dict(zip(labels.tolist(), counts.tolist())))这段检查代码逐个客户端打印样本数和类别计数参数client_idx是第 3 章切分函数的返回值。逻辑上它不参与训练只做数据体检如果某个客户端n小于batch_size要么调小 batch要么在本文还有配套的精品资源点击获取
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑