简介：面向希望理解横向联邦学习流程并动手做本地实验的Python开发者与学习者，该项目包以多客户端模拟为核心，帮助在无需共享原始数据的前提下完成本地训练、参数上传与服务器聚合的闭环，适合算法验证、课程演示和入门实践。资源共25个文件，主体为5个Python脚本，配合pyc缓存、xml与iml等IDE配置，以及CIFAR-10数据分片和元数据文件，压缩包约302.68MB，足以搭建可复现的小规模联邦训练环境。已有729人学习下载。包内围绕客户端训练、服务端聚合、模型定义、数据加载和配置入口组织代码，读者可据此对照实现联邦平均流程，观察多轮通信中的参数更新与全局模型变化，并进一步调整聚合策略、客户端数量或训练轮数，积累分布式机器学习的调试与实验经验。

1. 本地模拟横向联邦学习：一台笔记本跑通 FedAvg 的完整链路

想验证一个联邦学习想法，最省事的路径不是申请几台机器，而是在一台笔记本上把横向联邦学习的完整链路模拟出来：多个客户端各持一份分布不同的数据，本地训练几轮，只回传模型参数，服务端按样本量加权平均再广播回去。这套 Python 写的本地模拟框架干的就是这件事——它把“数据不出域”拆成能打断点、能打印张量的具体代码，适合刚接触联邦学习想跑通 FedAvg 的开发者，也适合要快速对比不同聚合策略的人。代价是慢，收益是每一行都能改。

2. 架构与目录：一次全局轮次里的六个角色

模拟框架最容易被写坏的地方，是把“分布式”三个字照搬进代码：开多进程、起 socket、搞消息队列，最后跑十分钟才出来一个 round，调试还得看日志猜。本地模拟的正确姿势是逻辑上分布式、物理上单进程串行——客户端之间不共享内存，但串行执行；通信开销不靠网络测量，靠参数量乘字节数算出来。这样一来，你可以直接给某个客户端的训练循环下断点，看它的 loss 是否发散，也可以在聚合前把每份权重打印出来比对，而不需要跨进程调试。

这套代码的整体数据流是单向的：simulate.py解析参数 →data.py切分数据并生成客户端 →server.py按轮次调度 → 每轮抽样若干客户端 →client.py深拷贝全局模型做本地训练 →aggregators.py把回传的 state_dict 加权合并 →metrics.py记录准确率和通信量。全程没有一处真正的并发，但每一步的语义都和在真实联邦场景里一致，这就够了。

2.1 角色与数据流：一个 round 里到底发生了什么

拆开看，一次全局轮次包含五步，顺序固定，缺一步都会导致结果不可复现。第一步是客户端采样，从 N 个客户端里按比例frac无放回抽一批，抽中的客户端集合决定这一轮的聚合权重分布；第二步是状态快照，把当前的全局state_dict深拷贝给每个被选中的客户端，注意必须是深拷贝，浅拷贝会让所有客户端共享同一份参数张量，本地训练时互相覆盖；第三步是本地多轮 SGD，每个客户端只在自己那份数据上跑local_epochs个 epoch；第四步是回传与聚合，服务端拿到的不是梯度而是完整权重，按各客户端样本量加权平均；第五步是评估与落盘，用统一测试集算准确率，同时把本轮参与客户端数、总样本数、上下行字节数写进指标文件。

这里有个容易被忽略的语义点：聚合权重用的是样本量而不是客户端数量的算术平均。当各客户端数据量差异很大时（Non-IID 切分几乎必然导致这种情况），算术平均会让只有几十条样本的客户端和小几千条样本的客户端话语权相同，收敛曲线会明显变差。所以client_update的返回值里除了 state_dict，还要带上len(loader.dataset)，这个数字会一路传到聚合函数里。

2.2 目录结构：八个文件撑起整套流程

代码量控制在千行以内，好处是你读一遍就能全改。目录按职责切分，不按“客户端/服务端”这种物理角色切，因为本地模拟里它们本来就是同一个进程里的两个函数。

fedtrain/ ├── configs/ │ └── default.yaml # 实验超参：轮次、客户端数、alpha、聚合器名 ├── fedtrain/ │ ├── __init__.py │ ├── data.py # 数据集加载 + Dirichlet 非IID切分 │ ├── client.py # 客户端本地训练，返回权重与样本量 │ ├── server.py # 轮次调度、客户端采样、评估 │ ├── aggregators.py # FedAvg / FedProx，带注册表 │ ├── models.py # 小 CNN 与两层 MLP，够跑 CIFAR/MNIST │ ├── metrics.py # 准确率、通信量统计、CSV 落盘 │ └── simulate.py # CLI 入口，只做参数解析和组装 ├── scripts/ │ ├── run_fedavg.sh │ └── run_fedprox.sh └── requirements.txt

依赖只有torch、numpy、pyyaml、tqdm四个，没有额外的联邦学习框架。这一点很关键：用现成框架跑模拟，出问题时你分不清是框架的调度逻辑在作怪还是自己的算法写错了；从零写一遍，deepcopy该放哪儿、随机种子该在哪一层设，全都暴露在明面上，debug 的成本低得多。models.py里的小 CNN 建议保持三层卷积加一层全连接，参数量在一百万上下，单机 CPU 跑 CIFAR-10 的一百轮大概十几分钟，够用来做对比实验，又不至于等到失去耐心。

3. 核心实现拆解：数据切分、本地训练与加权聚合

真正决定模拟结果可信度的，是三段代码：怎么切数据、怎么训本地、怎么合权重。其余都是外围。这三段代码里密密麻麻全是坑，写法上的细微差别会让曲线完全不同，所以下面逐段拆。

3.1 非独立同分布切分：Dirichlet 才是模拟的起点

如果所有客户端的数据都是均匀随机的，那联邦学习的难度就和集中式训练没区别了，测出来的曲线好看但没有参考价值。真实场景里每个客户端的数据分布都是偏的，模拟时用 Dirichlet 分布来造这种偏斜最方便，alpha越小越偏，alpha=0.1时极端情况下一个客户端可能只有两三个类别。

import numpy as np def dirichlet_split(targets, num_clients, alpha=0.5, seed=0): """按类别做 Dirichlet 采样，返回每个客户端的样本索引列表""" rng = np.random.default_rng(seed) targets = np.asarray(targets) num_classes = int(targets.max()) + 1 client_idx = [[] for _ in range(num_clients)] for c in range(num_classes): idx_c = np.where(targets == c)[0] rng.shuffle(idx_c) # 每个类别独立采一组比例，归一化后按比例切分该类的全部样本 props = rng.dirichlet(np.repeat(alpha, num_clients)) cuts = (np.cumsum(props) * len(idx_c)).astype(int)[:-1] for cid, part in enumerate(np.split(idx_c, cuts)): client_idx[cid].extend(part.tolist()) # 丢掉切分后为空的客户端，否则训练时会出现除零和空 loader return [idx for idx in client_idx if len(idx) > 0]

这段代码的逻辑是“按类切”而不是“按样本切”：先遍历每个类别，把该类样本打乱，再用一组 Dirichlet 比例决定这个类别分给各个客户端多少条。这样保证每个类别的样本都被完整分完，不会丢数据，同时客户端的类别分布天然偏斜。np.repeat(alpha, num_clients)是给每个客户端一个相同的浓度参数，让采样出来的比例向量和为 1；np.cumsum(props) * len(idx_c)转成累积位置，[:-1]去掉末尾那个必然等于总长的切点，剩下的就是np.split需要的分割位置。

参数上，alpha建议从 0.5 起步，做消融时再跑 0.1 和 1.0 三档；num_clients设 50 到 100 比较合适，太少体现不出采样效果，太多单机跑不动。最后那行过滤空客户端是血泪经验——alpha小于 0.1 时经常切出空列表，如果不剔除，后面len(loader.dataset)直接抛除零异常，报错信息还指不到切分函数上。

3.2 客户端本地训练：深拷贝和权重收集

本地训练函数只干两件事：在自己的数据上跑若干 epoch，然后把权重和样本量一起返回。写的时候有两个细节必须注意，一是不要修改传进来的模型对象，二是要把张量搬到 CPU 并 detach，否则显存会一路涨到 OOM。

import copy import torch def client_update(global_model, loader, args, device): """在单个客户端上做本地训练，返回 (权重字典, 样本量)""" # 每个客户端持有独立副本，避免本地更新污染全局模型 model = copy.deepcopy(global_model).to(device) optimizer = torch.optim.SGD( model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4 ) criterion = torch.nn.CrossEntropyLoss() model.train() for _ in range(args.local_epochs): for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() loss = criterion(model(x), y) loss.backward() optimizer.step() # 回传 CPU 上的副本，防止多客户端之间共享显存张量 state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} return state, len(loader.dataset)

copy.deepcopy这一行是整个模拟框架的正确性基石，漏掉它就会出现“客户端训练明明是独立的，结果第二轮开始所有客户端的初始权重都不一样”这种玄学现象——因为大家操作的是同一份参数，谁最后跑完谁说了算。detach().cpu().clone()三个操作缺一不可：detach断开计算图避免显存泄漏，cpu把张量挪出设备，clone保证返回值不和服务端的全局模型共享存储。

args.local_epochs一般设 1 到 5，设太大（比如 20）会让本地模型严重过拟合自己的偏斜数据，聚合之后全局模型反而倒退，这就是常说的客户端漂移。学习率用 SGD 加动量就够，别一上来就上 Adam，联邦场景下自适应优化器的聚合行为没那么直观，做对比实验时容易解释不清。weight_decay加一点点有助于抑制本地过拟合。

3.3 服务端聚合：加权平均与整型缓冲区的坑

FedAvg 的公式一行就能写完，但工程实现里最容易翻车的是 Buffer 层。BatchNorm 的num_batches_tracked是 int64 张量，直接参与浮点加权平均会类型不匹配报错；即使强行转 float 再转回来，语义上也说不通。

import torch def fedavg(global_state, client_states, weights): """按样本量加权平均，浮点参数求加权和，整型缓冲取首个客户端""" total = float(sum(weights)) agg = {} for k, v in global_state.items(): if v.dtype.is_floating_point: acc = torch.zeros_like(v, dtype=torch.float32) for state, w in zip(client_states, weights): acc += state[k].float() * (w / total) agg[k] = acc.to(v.dtype) else: # num_batches_tracked 这类计数器没有平均意义，直接沿用 agg[k] = client_states[0][k].clone() return agg

外层遍历的是global_state的键，保证聚合结果和全局模型的键集合严格一致，少一个键load_state_dict就会严格模式报错。v.dtype.is_floating_point是判断是否需要加权的那道门槛，浮点参数走加权求和，整型缓冲走直接拷贝。acc.to(v.dtype)那一步别省，某些模型权重是 float16，累加完不转回去会在加载时出现类型告警。

权重归一化放在循环内部做（w / total），比先累加再除更省内存，也避免了浮点累加误差被放大的问题。如果要做 FedProx，改动点就在这里：聚合公式不变，变化发生在客户端的损失函数里，加一项mu/2 * ||w - w_global||^2的近端项，mu取 0.001 到 0.01 之间，用来把本地模型往全局模型那边拽一点。这两种聚合器可以共用同一个fedavg函数，只是客户端的训练逻辑不同，这也是把聚合和训练拆成两个文件的原因。

3.4 轮次调度：把上面三段串起来

调度层代码不长，但它决定了实验能不能复现。核心是每轮重新采样客户端、串行调用、聚合、评估，并且在每轮开始前用固定的随机数生成器。

rng = np.random.default_rng(args.seed) for rnd in range(1, args.rounds + 1): # 每轮独立采样，frac 控制参与比例 m = max(1, int(args.frac * len(client_loaders))) selected = rng.choice(len(client_loaders), size=m, replace=False) states, weights = [], [] for cid in selected: state, n = client_update(global_model, client_loaders[cid], args, device) states.append(state) weights.append(n) global_model.load_state_dict(fedavg(global_model.state_dict(), states, weights)) acc = evaluate(global_model, test_loader, device) metrics.log(rnd, acc, m, sum(weights), param_bytes * 2 * m)

replace=False保证同一轮里不会重复抽到同一个客户端，max(1, ...)防止frac太小导致size=0时choice抛异常。通信量按“参数字节数 × 2 × 参与客户端数”估算，上行下行各算一次，param_bytes在模型构建后算一遍存成常量即可。这个数字看着粗糙，但用来做“通信轮次 vs 精度”的对比曲线完全够用，也是模拟框架最有价值的产出之一——它能量化“多跑十轮本地训练省下多少通信”。

4. 跑起来：从零到第一条收敛曲线的四步

代码能读和代码能跑是两件事。这一章按实际动手顺序走一遍，包括环境、配置、启动和结果判读，中间穿插几个我踩过的环境坑。

4.1 环境准备与依赖

Python 版本建议 3.9 以上，PyTorch 用 CPU 版就够，除非你要把device改成cuda加速。依赖只有四个，装完直接能跑。

python -m venv .venv && source .venv/bin/activate pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu pip install numpy pyyaml tqdm python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

第一行建虚拟环境，隔离掉系统里可能存在的旧版 torch。第二行装 CPU 版，体积小、装得快，做模拟实验完全够用。第三行补齐其余三个依赖。最后那行验证是必须的，输出里第二个值为False说明走的是 CPU 路径，如果你在代码里硬编码了.cuda()，这里就是第一个翻车点。数据集的下载路径建议通过环境变量指定到一个固定目录，别让torchvision默认往家目录塞，多跑几个实验后清理起来很烦。

4.2 配置项与启动命令

配置用 YAML 加命令行覆盖，两者合并优先级是命令行更高。几个关键参数的取值直接决定实验要跑多久。

启动命令走模块方式，避免相对导入报错：

python -m fedtrain.simulate \ --dataset cifar10 --alpha 0.3 --clients 50 --frac 0.2 \ --rounds 100 --local-epochs 2 --batch-size 64 \ --agg fedavg --lr 0.01 --seed 20240501 --out runs/fedavg_a03

参数里--alpha控制数据偏斜程度，--frac控制每轮参与的客户端比例，--agg在fedavg和fedprox之间切换，--seed是全局种子，会同时传给数据切分和客户端采样两个随机源。--out指定的目录下会生成metrics.csv和config.yaml的快照，后者是后悔药——两周后回头看你绝对记不清当时跑的是哪组参数，把配置一起存下来能省掉大量重跑。

4.3 日志与结果判读

跑起来之后终端会按轮打印进度条，每轮输出参与客户端数、聚合样本量、测试准确率和累计通信量。真正要看的是落盘的metrics.csv，用几行 pandas 就能画出对比曲线。

import pandas as pd import matplotlib.pyplot as plt df = pd.read_csv("runs/fedavg_a03/metrics.csv") # 平滑处理，单轮的抖动会掩盖整体趋势 df["acc_smooth"] = df["test_acc"].rolling(5, min_periods=1).mean() plt.plot(df["round"], df["acc_smooth"], label="FedAvg alpha=0.3") plt.xlabel("communication round") plt.ylabel("test accuracy") plt.legend() plt.savefig("runs/curve.png", dpi=150)

rolling(5)的窗口平滑几乎是必须的，原始曲线在 Non-IID 下抖动幅度能有五六个百分点，不平滑根本看不出趋势。判读时有三个信号值得盯：一是前二十轮的爬升速度，太慢通常是学习率或local_epochs设小了；二是中段是否出现明显回落，出现回落基本都是客户端漂移；三是最终精度与集中式训练基线的差距，alpha=0.3时差三到五个点属于正常范围，差十几个点就得回去查聚合权重是不是用了算术平均。

做对比实验时，务必保证除被比较的变量外其他全部一致，包括seed、客户端采样序列、数据切分结果。最省事的做法是先跑一次把切分好的索引存成.npz，后续所有实验都从同一个文件加载，这样数据集层面的随机性就被彻底固定住了。

5. 避坑与排查：本地模拟联邦学习最常见的五个翻车点

模拟框架的报错往往不指向真正的原因，因为一条链路上叠了数据切分、本地训练、聚合、评估四层。下面五条是我在不同项目里反复遇到的，按“现象 → 原因 → 解决”记下来，出问题时对着排查比读源码快得多。

5.1 数据与切分环节

现象：第一轮准确率就跑到 90% 以上，后面几乎不涨，曲线平得像一条直线。原因：切分函数写成了按样本随机划分，每个客户端拿到的都是全类别的均匀分布，等于把集中式训练拆成几份跑，联邦学习要模拟的异构性完全没体现。解决：检查切分是否按类别做 Dirichlet 采样，最简单的验证是打印每个客户端的类别计数，alpha=0.3时应该能看到明显的长尾分布，如果每个客户端的类别计数差不多，说明切分逻辑退化了。另外一个隐蔽原因是把测试集也一起切了，导致评估用了训练数据。

现象：训练中途抛ZeroDivisionError或者num_samples should be a positive integer。原因：alpha设得太小（比如 0.05）时，某些类别在某些客户端上分到的样本数是 0，np.split会产生空数组，最后拼出来的客户端索引列表是空的。解决：在切分函数末尾过滤掉空客户端，同时打印被丢弃的数量；如果丢弃比例超过 10%，说明alpha或客户端数设置不合理，应该调大alpha或者减少客户端数，而不是硬过滤。

5.2 训练与聚合环节

现象：第二轮开始，每个客户端的初始权重都不一样，或者准确率突然掉到随机猜测水平。原因：client_update里没有深拷贝模型，所有客户端共用同一份参数张量，后跑的覆盖先跑的；或者聚合时把client_states里的张量直接赋给了全局模型，导致后续本地训练修改了全局权重。解决：确认copy.deepcopy在client_update内部、load_state_dict在聚合之后；调试时可以在聚合前后各打印一次全局权重的哈希值，正常情况下聚合前每轮的初始哈希应该相同。

现象：load_state_dict报Missing key(s) in state_dict: ...num_batches_tracked。原因：聚合函数对整型缓冲区也做了加权平均，或者干脆跳过了非浮点键。解决：按 dtype 分支处理，浮点参数加权、整型缓冲取首个客户端，并且外层遍历以global_state的键为准，这样能保证键集合严格对齐。还有个小概率情况是模型里有自定义 buffer 未注册，需要检查模型定义里的register_buffer。

现象：跑到三四十轮显存或内存持续上涨，最后 OOM。原因：本地训练返回的张量没有detach，保留了计算图，每轮累积一点，几十轮下来就撑爆了。解决：返回值统一走v.detach().cpu().clone()；另外注意评估函数要用torch.no_grad()包住，否则测试集过一遍也会建图。这两个地方都是不会报错只会慢慢变卡的隐性泄漏，用nvidia-smi或者psutil盯着内存曲线最容易发现。

现象：同样的参数跑两次结果差好几个点。原因：随机种子只设了torch.manual_seed，没设numpy的生成器；或者 DataLoader 的shuffle=True用了全局随机状态，多线程下不可复现。解决：数据切分用独立的np.random.default_rng(seed)，客户端采样用另一个独立的生成器，DataLoader 加worker_init_fn或者直接num_workers=0。本地模拟场景下num_workers=0的速度损失可以接受，换来的是彻底的可复现性，这笔账很划算。

6. 二次开发：把模拟框架改成你自己的算法试验台

想在这套代码上加自己的聚合算法，改动点其实只有一个：aggregators.py里的注册表。我一般保持fedavg不动作为基线，新算法写成独立类，通过--agg切换，这样任何一次实验都能随时退回基线做对照。接入方式像下面这样，注册装饰器负责把名字映射到类，simulate.py里按字符串取用即可。

from fedtrain.registry import AGGREGATORS import torch @AGGREGATORS.register("trimmed_mean") class TrimmedMean: """去掉最大最小各 k 份后取平均，用来抵抗个别客户端的异常权重""" def __init__(self, k=2): self.k = k def aggregate(self, global_state, client_states, weights): agg = {} for key, ref in global_state.items(): if not ref.dtype.is_floating_point: agg[key] = client_states[0][key].clone() continue stacked = torch.stack([s[key].float() for s in client_states]) # 按客户端维度排序后裁掉两端，再沿客户端维度求均值 sorted_v, _ = torch.sort(stacked, dim=0) kept = sorted_v[self.k: stacked.shape[0] - self.k] agg[key] = kept.mean(dim=0).to(ref.dtype) return agg

关键在于接口统一：入参永远是(global_state, client_states, weights)，出参永远是与global_state键集合一致的字典，这样任何聚合器都能被调度层无差别调用。weights在这个例子里没用上，但接口保留着，因为大多数加权类算法都依赖它。写完别忘了在配置里加一行默认参数，否则__init__的参数只能靠命令行透传，用起来别扭。

验证新算法时我有个固定流程：先把seed固定，跑一遍基线记录每轮准确率，再跑新算法，逐轮对比差值曲线；如果新算法在前期领先但后期被反超，多半是聚合权重的问题而不是算法本身；如果两条曲线从第一轮就分叉，去查是不是聚合后没有正确load_state_dict。做完这三步再去看最终精度，能省掉大量误判。

从那以后我每次新加一个聚合器，都强制先把客户端采样序列和切分索引 dump 到文件里复用，再谈效果对比——不然两次实验的随机性叠在一起，你根本说不清提升是算法带来的还是抽样运气好。希望帮到你。

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