简介一份基于Python的横向联邦学习本地模拟工程面向想学习联邦学习原理、分布式训练思路的开发者与学生。工程以server.py与client.py为运行核心配合models.py、datasets.py、main.py等源码以及conf.json参数配置和CIFAR-10数据集完整呈现多客户端本地训练、模型更新上传、服务器平均聚合与全局模型下发的循环。压缩包共25个文件总大小约302.68MB文件类型以Python脚本为主另有pyc编译缓存、PyCharm工程XML配置、CIFAR-10分批数据文件目录结构清晰适合直接打开运行或按模块研读。已有725人学习/下载。对希望快速上手联邦学习实验、完成课程设计或搭建本地模拟环境的读者来说这份资源能提供可直接运行的参考代码和完整数据支撑。1. 从一台机器跑出联邦学习这个模拟项目到底在解决什么想验证联邦聚合算法但没有分布式集群云上临时租机器成本又高怎么办某开发者的答案是把每个客户端模拟成内存里的一个独立对象让服务器在本地完成参数分发、训练和聚合。这个项目就是用Python实现的本地横向联邦学习模拟器基于FedAvg算法支持自定义客户端数量、数据分布和模型结构不需要分布式框架和网络通信普通CPU就能跑完联邦学习全流程。适合算法工程师验证聚合策略、学生复现论文以及所有想在一台机器上理解联邦学习原理的开发者。项目核心逻辑不到两百行依赖只有PyTorch、numpy和tqdm拿到源码后改改配置文件就能跑出自己的实验曲线。相比部署真实联邦系统这个模拟器把注意力集中在算法本身非常适合用来对比IID和Non-IID数据下的收敛差异或者试验各种聚合策略。2. 项目架构与目录结构模拟联邦的骨架怎么搭2.1 整体分层数据层、客户端层、服务器层、聚合层这个模拟器没有真实网络通信层因为所有客户端都运行在同一个进程里。为了贴近真实横向联邦学习它把整个流程拆成四个逻辑层数据层、客户端层、服务器层、聚合层。数据层负责加载原始数据集并按配置切成多个子集每个子集对应一个客户端。这里支持两种切分模式IID模式把样本随机打乱后均匀分配Non-IID模式按样本类别分组再把不同类别组合分配给不同客户端。Non-IID是联邦学习里最重要的实验条件因为真实场景中每个参与方的数据分布几乎都是异构的如果所有客户端数据都是均匀的聚合策略很难暴露出问题。客户端层是一个Python类每个实例持有自己的本地数据、模型参数和训练超参数。它的职责很单一接收服务器下发的全局模型参数用它初始化本地模型在自己的数据上训练若干个epoch最后把更新后的参数和本地样本数返回给服务器。注意整个训练过程中服务器和客户端之间传递的是模型参数不是原始数据这保证了模拟过程符合联邦学习的隐私前提。服务器层负责协调所有客户端。它维护一个全局模型每一轮通信开始时把全局参数广播给本轮参与训练的客户端等所有客户端返回更新后调用聚合层计算新的全局参数然后进入下一轮。服务器不接触任何本地训练数据只接触参数和样本数量这和真实联邦学习的规则一致。聚合层实现了FedAvg的核心公式。这个层只有几十行代码但它是整个项目最容易写错的地方聚合时每个客户端的权重必须用本地样本数占总样本数的比例来计算而不是简单的算术平均。后面我们会看到这个加权系数直接影响收敛效果写错时程序不会报错只会默默地让精度上不去是典型的“黑匣子”问题。这四个层的划分本质上对应了设计模拟器时要考虑的三个问题如何构造异构数据、如何实现本地训练、如何做参数聚合。理解了这三件事后面读代码就会很顺手。2.2 目录结构逐层拆解从数据切分到聚合输出项目文件组织得很清晰拿到手先花两分钟过一遍目录后面改起来就不会迷路。fed_sim/ ├── config.py # 全局配置路径、超参数、模拟参数 ├── data.py # 数据加载与IID/Non-IID切分 ├── model.py # 模型定义一个两层全连接网络 ├── client.py # 客户端类本地训练与返回参数 ├── server.py # 服务器类参数分发、聚合、评估 ├── utils.py # 工具函数日志、指标计算、种子设置 ├── main.py # 入口解析命令行参数启动模拟 └── requirements.txt # 依赖清单config.py 是个新手友好的设计。所有实验参数都集中在这里比如客户端数量、每轮参与比例、本地epoch数、batch_size、学习率、通信轮数以及是否启用Non-IID切分。改实验条件时不用去翻各个模块的源码打开这个文件改一行就能生效。data.py 负责把原始数据集变成每个客户端的私有数据。它内部维护一个字典键是客户端编号值是该客户端的数据loader。切分逻辑写在两个函数里split_iid()和split_noniid()。前者用随机索引均匀分配后者先按标签分组再按分组组合分配给客户端。关于Non-IID切分下一节会展开讲。client.py 和 server.py 是对应前面说的两个核心角色。client 里的set_parameters()和get_parameters()负责参数转换train()负责本地训练server 里的aggregate()调用聚合公式产生新的全局参数。这两个类都不需要感知对方在内存中如何存在接口上已经做了抽象。model.py 定义了一个简单的两层全连接网络输入维度跟数据集特征数对齐隐藏层128维输出为类别数。这个模型可以整体替换成CNN或Transformer只要保证get_parameters()和set_parameters()两个接口存在模拟器就能继续跑。main.py 是整个项目的调度中心。它读取配置创建数据和客户端初始化服务器然后进入通信循环。每一轮循环会打印训练精度、损失和聚合耗时方便观察收敛过程。2.3 数据切分模拟Non-IID切分如何贴近真实场景Non-IID切分是联邦学习模拟中的关键操作。如果所有客户端的数据分布都一样聚合就会退化成普通的批量训练测不出聚合算法的价值。下面这段代码展示data.py中Non-IID切分的核心思路def split_noniid(dataset, num_clients, num_shards200): # 先按类别把样本索引分组 labels [y for _, y in dataset] label_to_indices {} for idx, label in enumerate(labels): label_to_indices.setdefault(label, []).append(idx) # 把每个类别的样本平均切成若干shard shards [] for label, indices in label_to_indices.items(): random.shuffle(indices) shard_size len(indices) // (num_shards // len(label_to_indices)) for i in range(0, len(indices), shard_size): shard indices[i:i shard_size] if len(shard) shard_size: shards.append(shard) # 每个客户端随机分配2~3个shard random.shuffle(shards) client_data {i: [] for i in range(num_clients)} for cid in range(num_clients): for _ in range(random.randint(2, 3)): if shards: client_data[cid].extend(shards.pop()) return client_data这段代码的逻辑是先把整个数据集按标签分组再把每个标签内的样本切成大小相等的分片shard最后让每个客户端随机拿几个分片。这样不同客户端拿到的类别组成可能完全不同有的客户端只有数字0和1有的只有7和9模拟出了真实场景中数据偏斜的现象。参数num_shards控制分片总数它决定了数据切得有多碎。分片越多每个分片包含的类别越单一客户端之间的数据分布差异就越大。这里默认200个分片对10类数据集来说每个类被切成20个分片每个客户端随机拿2~3个既能保证每个客户端有一定数量样本又能产生足够明显的分布偏移。如果希望模拟更极端的数据异构可以把num_shards调小到80或50客户端拿到的类别就会越少。这里有个容易踩坑的细节如果num_shards设置得比客户端数量还小分片会被抽完后面的客户端会拿到空数据训练时会报“data loader is empty”的错误。所以代码里加了if shards的判断但在真实运行时最好在配置文件里检查num_shards是否大于等于num_clients * 2否则部分客户端可能样本太少导致模型训不起来。IID切分就简单得多直接把样本索引随机打乱平均分成num_clients份每个客户端拿到类别比例与整体一致。两种模式共用一个数据加载接口在config里切换一个布尔值就行。3. 核心代码拆解FedAvg从客户端到服务器的完整链路3.1 客户端本地训练模型初始化与更新客户端的核心职责是在自己的私有数据上训练模型。为了让代码尽量通用客户端不感知神经网络的具体结构只负责参数的存取和梯度更新。下面是client.py中客户端类的核心实现class Client: def __init__(self, cid, model_fn, data_loader, config): self.cid cid self.model model_fn() # 创建一个全新的本地模型 self.data_loader data_loader self.optimizer torch.optim.SGD(self.model.parameters(), lrconfig.lr) self.local_epochs config.local_epochs self.loss_fn torch.nn.CrossEntropyLoss() def set_parameters(self, params): # 用服务器下发的全局参数覆盖本地模型 for local_param, global_param in zip(self.model.parameters(), params): local_param.data.copy_(global_param) def get_parameters(self): return [param.data.clone() for param in self.model.parameters()] def train(self): self.model.train() total_loss 0 for epoch in range(self.local_epochs): for x, y in self.data_loader: self.optimizer.zero_grad() output self.model(x) loss self.loss_fn(output, y) loss.backward() self.optimizer.step() total_loss loss.item() return self.get_parameters(), len(self.data_loader.dataset)这里set_parameters用data.copy_而不是直接赋值是为了保证本地模型的参数张量与全局模型参数张量共享存储位置。如果直接赋值新的张量会断开与优化器的连接导致后面optimizer.step()更新不到任何参数。这个细节处理不好联邦学习会变成“只聚合不训练”而且不会报错属于典型的翻车场景。train方法返回两个值更新后的参数列表和本地样本数。样本数用于后续FedAvg加权len(self.data_loader.dataset)在DataLoader的drop_lastFalse下等于该客户端持有的总样本数。注意这里返回的参数是通过clone()克隆出来的避免返回了还在被优化器引用的临时张量。虽然PyTorch的梯度释放机制通常不会出问题但克隆一份能让聚合阶段的内存管理更干净。优化器选用的是SGD而不是Adam。这是联邦学习实验中的常见选择。因为FedAvg论文里标准配置就是SGDmomentum很多后续研究也沿用这个基础设定。如果你想在本地模拟时用Adam可以直接在__init__里把torch.optim.SGD替换成torch.optim.Adam然后观察聚合收敛速度是否变化。不过要注意Adam的动量状态是每个客户端自己维护的不同客户端的学习率行为差异会让聚合后的全局模型产生额外扰动。本地epoch数的选择也很微妙。local_epochs1时每个客户端每轮只在自己的数据上做一遍前向和反向更新幅度小聚合稳定性高但收敛速度慢。把local_epochs调到5或10本地模型会更充分地拟合私有数据每轮通信带来的信息量更大但也可能让模型偏移到局部最优。在Non-IID场景下过大的本地epoch会让聚合后的全局模型精度大幅震荡这是联邦学习领域里一个经典的“本地更新累积漂移”问题。3.2 服务器聚合FedAvg加权平均的实现聚合是联邦学习的核心步骤。FedAvg的思想很简单全局参数的更新等于所有客户端参数更新的加权平均权重是该客户端样本数占总样本数的比例。下面server.py中的聚合函数def fed_average(global_params, client_updates, client_samples): # client_updates: list of list of torch.Tensor # client_samples: list of int total_samples sum(client_samples) # 按样本数比例计算权重 weights [num / total_samples for num in client_samples] # 初始化聚合参数为0 aggregated [torch.zeros_like(param) for param in global_params] for update, weight in zip(client_updates, weights): for param, grad in zip(aggregated, update): param.add_(grad * weight) return aggregated这段代码最容易被忽视的是权重计算。如果简单地用1 / num_clients做算术平均在客户端数据量差异较大的场景下小样本客户端的更新会被大样本客户端淹没收敛效果明显变差。FedAvg的加权设计就是为了解决这个数据量不均衡问题。另一个值得注意的点是torch.zeros_like(param)。这里必须使用与全局参数相同形状和设备的零张量否则后面add_会因为设备或形状不一致报错。在本地模拟场景中所有客户端都在内存里设备都是CPU但如果你把代码移植到多卡环境这里就需要显式指定设备。实现细节上为什么使用param.add_(grad * weight)而不是直接赋值因为grad * weight会创建一个新的张量多次累加到aggregated上能够避免形状意外变化。如果某个客户端的模型结构定义不同这里会立刻暴露形状不一致的问题比起直接赋值更安全。聚合后服务端一般会做一次全局模型评估。评估用的是服务器持有的测试集这个测试集不参与任何客户端训练是模拟器独立准备的。评估过程很简单把模型设为eval()模式计算在测试集上的精度。值得注意的是服务端的全局模型和客户端本地模型结构相同但评估时不需要复制到客户端直接在服务端调用模型即可。3.3 轮次调度与模型同步一次通信的完整流程main.py中的主循环模拟了多个通信轮次。每一轮就是一次完整的“广播-训练-聚合”过程def run_simulation(server, clients, rounds, participation_ratio1.0): for round_id in range(1, rounds 1): # 1. 服务器广播全局参数 global_params server.get_global_params() # 2. 随机选择参与本轮训练的客户端 num_active max(1, int(len(clients) * participation_ratio)) active_clients random.sample(clients, num_active) # 3. 每个客户端在本地训练返回更新和样本数 updates, samples [], [] for client in active_clients: client.set_parameters(global_params) update, sample_count client.train() updates.append(update) samples.append(sample_count) # 4. 服务器聚合更新全局模型 new_params fed_average(global_params, updates, samples) server.set_global_params(new_params) # 5. 评估并记录日志 acc server.evaluate(test_loader) print(fRound {round_id}, Accuracy: {acc:.4f})participation_ratio是一个很有价值的模拟参数。真实联邦学习的参与方不可能每轮都全部在线通常只有一部分客户端响应。这个参数控制每轮随机挑选的客户端比例。当设置为0.5时每轮只有一半客户端参与训练可以观察“部分参与”对收敛速度的影响。如果设置为1.0就是标准的全量参与。这里模拟的“通信”其实就是同步调用客户端对象的train()方法。因为在单进程中客户端训练会顺序执行这等价于同步联邦学习——服务器必须等所有选中的客户端完成训练后才进行聚合。如果后续想改成异步联邦只需要把这段循环改成线程池提交并在每个客户端完成时立即聚合一次。主循环里还隐藏着一个评估时机的细节测试集评估应该放在聚合完成之后而不是聚合之前。为了观察每一轮通信的效果评估必须在全局参数更新后执行。有些初学者会把评估放在广播之前那评估的其实是上一轮的模型导致日志中的精度比实际滞后一轮对比实验时容易产生困惑。4. 本地跑通与参数调优从命令行到实验报告4.1 环境准备与依赖安装项目依赖非常轻requirements.txt 里只有三个包torch1.10.0 numpy1.21 tqdm4.60安装用一行命令pip install -r requirements.txt建议在虚拟环境中安装。用conda创建Python 3.8环境然后执行上述命令。如果网络慢可以设置国内镜像源。另外如果机器有CUDA GPUPyTorch会默认使用GPU但本地模拟场景下客户端数量较多每个客户端模型很小CPU反而更适合因为GPU在大量小模型切换时开销远大于收益。项目默认在CPU上运行不需要额外配置。如果你用的是Apple Silicon芯片PyTorch 2.0以上版本支持MPS后端可以在config.py里加一行device mps来加速训练。但要注意torch.zeros_like在MPS上偶尔会有兼容问题遇到报错就切回CPU。4.2 运行命令与参数配置项目通过main.py接收命令行参数大部分参数都有默认值可以直接运行python main.py --clients 10 --rounds 20 --local_epochs 2 --batch_size 32 --lr 0.01 --noniid常用参数如下表参数默认值含义--clients10模拟客户端总数--rounds10全局通信轮数--local_epochs1每轮通信内本地训练epoch数--batch_size32本地训练batch大小--lr0.01本地SGD学习率--noniid无启用Non-IID数据切分默认IID--participation1.0每轮参与客户端比例从这些参数能看出实验可以调节的维度很多。--local_epochs控制本地计算量对应联邦学习中的“多步本地更新”策略值越大每轮通信减少的训练损失越多但也更容易产生客户端漂移。--lr在聚合收敛中扮演关键角色如果聚合后精度震荡一般先调低学习率。--clients决定模拟规模超过50个客户端时建议配合--participation使用否则单轮训练时间会线性增长。这里有一个实验设计上的建议先把--clients固定为10--local_epochs固定为1跑通IID模式作为基线。然后切换成Non-IID保持其他参数不变观察精度下降幅度。这个对比是理解联邦学习数据异构影响的最直观方式。如果直接在一个复杂参数组合上跑出了问题时很难判断是聚合算法的问题还是数据切分的问题。4.3 结果输出与指标解析程序每轮会输出类似下面的日志Round 1, Accuracy: 42.10%, Loss: 1.8923 Round 2, Accuracy: 51.33%, Loss: 1.6540 ... Round 20, Accuracy: 86.25%, Loss: 0.4311Accuracy是在服务器端用测试集评估全局模型得到的结果Loss是所有参与客户端本地损失的平均值。这两个指标的含义不一样全局精度反映模型收敛效果适合对比不同聚合策略本地平均损失反映客户端内部拟合程度如果它一直下降但全局精度不涨说明出现了客户端过拟合或数据分布偏移的问题。建议每跑完一组实验把日志保存下来画成“轮次-精度”曲线。项目没有内置绘图模块但日志格式是纯文本用pandas或Excel都能快速导入。我一般会额外加一个--log_path参数把每一轮的结果写入CSV方便后续对比。更严谨的做法是同时记录每个客户端样本数、每轮参与客户端编号。这些信息在实验复现时非常有用。因为随机种子一旦固定重复运行会得到完全一致的结果如果不固定每次运行参与客户端顺序不同精度曲线会有轻微抖动但这在合理范围内。4.4 常见失败场景的处理刚拿到项目时最容易遇到的错误是ModuleNotFoundError: No module named torch。这个原因很直接就是没有安装依赖或者安装到了错误的环境中。解决方法是确认虚拟环境是激活状态再执行pip list检查torch是否在列表里。第二个常见错误是参数形状不匹配导致聚合时报错。这通常发生在修改模型结构之后比如你改了隐藏层维度但之前训练保存的全局模型参数还是旧维度。解决方法是删除可能存在的checkpoint.pkl或者确保从头运行而不是加载旧模型。第三个是运行速度慢。如果客户端个数大于20participation_ratio为1.0每轮训练所有客户端总耗时可能达到几分钟。这不是死循环项目在启动时会打印预计参与客户端数量你可以中途用CtrlC停止然后调高--batch_size或调低--local_epochs来缩短时间。第四个问题来自数据集下载。项目默认使用某个内置的10类图像数据集首次运行时会自动下载。如果网络不通会自动寻找本地缓存的.npz文件。建议提前下载好数据集放到data/目录下避免每次实验卡在下载环节。第五个问题是精度曲线完全没有上升一直停留在初始值附近。这种情况先检查数据切分如果num_shards设置过大客户端可能没有拿到任何标签完整的样本或者样本数太少导致本地训练根本学不到东西。把数据分布打印出来看一遍往往能发现端倪。5. 避坑指南本地模拟联邦学习的血泪经验5.1 坑一Non-IID切分导致某些客户端数据类别缺失现象在Non-IID模式下训练全局精度始终在20%左右徘徊远低于IID模式的80%以上。查看每个客户端的类别分布发现某些客户端只包含2个类别甚至只有1个类别。原因这是Non-IID切分本身造成的类别不平衡。每个客户端拿到的shard类别不全本地模型只能学到少数类别的特征聚合时这些偏斜更新互相干扰导致全局模型收敛困难。解决先把num_shards调大比如从200调到400让每个客户端拿到更多但更碎的数据分片。或者每个客户端最少分配3~4个shard确保至少覆盖3个类别。如果仍然收敛差可以尝试在聚合时增加客户端的本地epochs让本地模型更好地拟合自己有限的数据分布。5.2 坑二全局模型初始化的传播问题现象修改了model.py中的网络结构后所有客户端训练出来的损失一直不下降但服务器端不会报任何错误。原因新建客户端时每个客户端都调用了model_fn()创建独立模型这些模型的参数初始化是随机的。服务器初始化的全局模型也是随机参数。第一轮广播时每个客户端用全局参数覆盖自己的参数这一步没问题。问题在于如果main.py在创建客户端之后才初始化服务器那么服务器参数永远是第一轮之前的随机值聚合时不会更新。解决检查主流程确保服务器初始化在创建客户端之后并且在模拟开始前把全局参数广播给所有客户端一次。我惯用的做法是在run_simulation的开头加一个断言检查服务器全局参数的第一个张量是否与客户端get_parameters()返回的形状一致不一致立刻报错。5.3 坑三聚合时权重顺序不一致现象聚合后精度比聚合前还低而且每轮精度变化非常不稳定甚至出现NaN空值。原因client_updates列表的顺序和client_samples列表的顺序不一致。比如在3.3节的循环中如果先收集samples再收集updates或者某个客户端训练失败被跳过就会导致权重对错位置。解决在循环中同步append两个列表确保索引一致。更稳妥的方式是让客户端train()返回(cid, params, count)三元组按cid排序后再喂给聚合函数。这样即使客户端执行顺序发生变化权重也能匹配。5.4 坑四模拟通信开销导致CPU占用过高现象客户端数量调到50后程序运行时间成倍增长CPU占用接近100%看着像是死循环。原因单个进程内顺序训练所有客户端每个客户端都要做前向、反向和参数克隆操作大量小张量频繁创建销毁CPU忙个不停。解决不需要真的上分布式。可以先用participation_ratio0.3减少每轮参与数再增大batch_size到64或128。如果还不够把客户端模型换成更小的结构隐藏层从128降到64。本地模拟的目的是验证算法不是比拼训练速度所以不需要追求所有客户端同时训练。5.5 坑五随机种子没有固定导致实验不可复现现象两次跑相同参数得到的结果曲线完全不一样甚至一次收敛一次发散。原因项目里使用了random、torch等模块的随机数生成器如果不设置种子数据切分、客户端选择、模型初始化都会变化。解决在main.py入口处调用utils.set_seed(42)。函数内部把Python的random、NumPy和PyTorch的随机种子都固定下来并把torch.backends.cudnn.deterministic设为True。这样每次运行的结果完全可复现做实验对比时才可信。6. 二次开发与改造思路把模拟器变成你的实验平台6.1 替换数据集与模型结构想把模拟器用在其他数据集上只需修改data.py和model.py。数据加载部分把内置数据集的读取换成自己的数据类保证返回的dataset支持索引访问和__len__同时标签是整数就能继续用原来的切分逻辑。模型方面只要保留get_parameters和set_parameters两个方法把内部网络换成CNN或简单Transformer聚合层不需要改动因为聚合只操作参数张量列表。例如把两层全连接网络换成LeNet风格的小卷积网络输入通道和图像尺寸需要匹配。我一般会把model_fn定义在model.py里通过config的参数控制隐藏层维度或卷积通道数这样实验时不用改代码只调配置。6.2 加入差分隐私噪声联邦学习经常需要加入隐私保护机制。在本地模拟器里加差分隐私非常方便修改点只在聚合阶段。FedAvg的隐私版本通常会在客户端上传参数前对参数加上高斯噪声噪声尺度由隐私预算决定。具体实现是在client.train()返回参数前遍历每个参数张量加上torch.randn_like(param) * noise_scale。噪声越大隐私保护越强但模型精度会下降。改造时要小心收敛性和评估指标。我在做这个实验时会同时记录噪声标准差观察不同噪声水平下的精度损失曲线。这个模拟器是进行这类对比实验的理想平台因为不需要考虑真实通信链路的加密开销。6.3 从同步联邦改成异步联邦同步联邦是每一轮等所有客户端完成异步联邦允许客户端在训练完成后立即上传更新。改动思路在3.3节已经提到用concurrent.futures.ThreadPoolExecutor替代串行循环每个客户端训练完成后立刻调用聚合函数更新全局模型不需要等待其他客户端。异步聚合时要小心参数版本问题客户端拿到的全局参数可能不是最新版本。可以在客户端对象上记录参数版本号聚合时用客户端返回的版本号与当前全局版本号比较如果落后太多就丢弃这次更新。这是异步联邦学习里常用的“陈旧性”处理。改造完成后你会看到精度曲线比同步模式更早进入平台期但波动也会更大。做异步改造时建议先在participation_ratio0.3的小规模场景下验证观察线程池调度是否稳定。我第一次改异步时就遇到过两个客户端同时聚合导致全局参数被覆盖的问题后来加了一把线程锁聚合操作同步化才解决。从那以后每次改造并发相关的逻辑我都会强制检查数据竞争和参数版本这两个点几乎能规避掉所有异步联邦的翻车事故。希望这个模拟器的踩坑记录能帮你少走一段弯路。本文还有配套的精品资源点击获取