简介一份用Python在本地模拟横向联邦学习完整流程的工程代码包面向希望通过实际项目理解联邦学习原理的机器学习开发者与学生。项目实现了多客户端本地训练、参数上传、服务端聚合和全局模型更新的基本闭环包含server、client、models、datasets等核心Python模块并配有配置文件与CIFAR-10数据集可直接运行观察训练效果也便于分析每一轮通信中的参数变化。资源共25个文件以.py源码为主附带pyc缓存、iml工程文件、xml配置等辅助内容以及多批次CIFAR-10图像数据压缩包约302.68MB已有725人学习。借助这份代码可省去自行搭建通信与数据预处理的时间把注意力集中在联邦学习机制本身模块划分清晰适合作为课程设计、毕业设计或入门实践项目的参考蓝本在此基础上调整数据划分、聚合策略或模型结构即可开展更多横向联邦学习实验。1. 本地模拟横向联邦学习一套代码把论文里的分布式训练跑明白先抛一个反直觉的结论理解联邦学习最快的路径不是搭多机多卡集群而是在一台电脑上用 Python 模拟出多个“虚拟客户端”。很多开发者在接触横向联邦学习时第一反应是找几台机器部署通信框架结果光环境配置就耗掉一周最后跑通的还是一个玩具 Demo。这个模拟项目解决的问题恰恰就在这里——它把横向联邦学习的“数据分布—本地训练—参数上传—全局聚合—下发更新”完整闭环压缩到单机环境里用进程或线程模拟多个参与方让你在不开任何分布式集群的前提下把 FedAvg 这类经典聚合算法的每一行代码都吃透。这个资源适合三类人正在做联邦学习课程设计或毕业设计的学生需要快速验证聚合算法思路的算法工程师以及想在简历里写出“熟悉联邦学习训练流程”但不想被运维细节劝退的后端开发者。它基于 Python 生态实现配合 PyTorch 或 TensorFlow 做底层模型训练所有代码逻辑都围绕横向联邦最核心的“数据不动模型动”展开。接下来我会从技术选型、目录结构、核心代码到运行方式完整拆解这套本地模拟实现让你拿到源码后能直接跑起来并且知道从哪一行开始改成你自己的算法。2. 技术选型与整体架构为什么选择进程模拟而不是真正的分布式通信2.1 技术栈定位Flower 框架 PyTorch 后端 数据集切分这套项目选型非常克制核心依赖只有三块。首先是深度学习框架 PyTorch模型定义、本地训练和梯度计算都基于它完成其次是联邦学习框架 Flower用于管理客户端注册、参数聚合和通信轮次最后是数据处理层项目内置了 MNIST 和 CIFAR-10 两个标准数据集的加载与切分工具。之所以选 Flower 而不是自己手写 Socket 通信是因为 Flower 天然支持在单机模式下启动多个客户端实例同时保持了与真实分布式部署一致的接口抽象这意味着你今天用模拟方式跑通的代码明天可以直接平移到多机环境。要理解这个项目的架构先要抓住横向联邦学习的一个本质特征参与方的数据特征维度相同但样本 ID 不同。在本地模拟中项目通过把完整训练集按样本维度切分成多个不重叠的分片来模拟这种数据分布。比如把 MNIST 的 60000 张图片分成 10 份每个客户端只持有其中 6000 张且这些图片的类别分布可能是不均匀的——这种不均匀正是联邦学习要解决的现实问题。2.2 核心数据流一个训练轮次的完整生命周期整个系统的数据流可以拆成四个阶段。第一阶段是服务器初始化创建全局模型并随机初始化参数第二阶段是客户端选择与下发服务器从注册的客户端列表中选择一部分参与本轮训练把当前全局模型参数下发第三阶段是本地训练与上传每个客户端用自己的本地数据训练若干轮然后把更新后的参数或参数差值上传回服务器第四阶段是聚合与广播服务器用 FedAvg 算法对收到的参数做加权平均更新全局模型然后进入下一轮。这个项目在代码层面把上述流程映射成了三个核心模块server.py负责聚合逻辑client.py负责本地训练utils.py负责数据切分和模型定义。你不需要理解 Flower 底层的网络通信细节只需要知道它提供了一个fl.server.start_server入口和fl.client.start_numpy_client入口前者启动聚合服务后者让每个客户端进程接入。项目用多进程方式启动多个客户端每个进程独立加载自己的数据分片模拟出“多个机构各自持有数据”的效果。2.3 目录结构速览每个文件承担什么职责federated-learning-simulation/ ├── server.py # 联邦服务器聚合策略、轮次控制、模型评估 ├── client.py # 联邦客户端本地训练、参数上传、模型下载 ├── utils.py # 工具函数数据集切分、模型定义、参数序列化 ├── run_simulation.py # 入口脚本配置参数、启动服务器和多个客户端进程 ├── config.yaml # 配置项客户端数量、本地训练轮数、聚合轮数、学习率等 └── README.md # 运行说明和环境依赖utils.py是整个项目的地基它定义了要训练的模型结构——通常是两层卷积加全连接层的简单 CNN以及数据切分的具体实现。如果你要换数据集或改模型结构绝大多数改动都集中在这个文件里。config.yaml则集中管理所有可调参数从客户端总数到每轮参与比例从本地 epoch 到学习率全部声明式配置改参数不需要动代码。这套目录设计的好处是职责单一初学者能快速定位每个文件的作用二次开发者也能在半小时内找到自己要改的位置。3. 核心模块拆解一个训练轮次从进入到返回的完整链路3.1 数据切分模块如何制造“非独立同分布”数据联邦学习模拟成败的关键在于数据切分是否真实。如果每个客户端拿到的数据都是随机均匀抽样那么各客户端的本地数据分布和全局分布几乎一致联邦学习的效果会和集中训练没有区别这就不符合真实场景。现实中各机构的数据分布往往是有偏的比如某家医院的病人数据偏老龄化另一家偏年轻化。# utils.py: 非独立同分布数据切分 import numpy as np from torch.utils.data import Subset def partition_data(dataset, num_clients, non_iid_degree0.5): 将数据集按非独立同分布方式切分给多个客户端 :param dataset: 完整数据集torchvision Dataset对象 :param num_clients: 参与联邦训练的客户端数量 :param non_iid_degree: 非独立同分布程度0表示完全随机1表示纯按类别划分 :return: client_data_indices: list每个元素是该客户端对应的样本索引列表 labels np.array([dataset[i][1] for i in range(len(dataset))]) class_indices {} for idx, label in enumerate(labels): class_indices.setdefault(label, []).append(idx) client_data_indices [[] for _ in range(num_clients)] for label, indices in class_indices.items(): np.random.shuffle(indices) # 将每个类别的样本按比例分给各客户端 if non_iid_degree 0.5: # 近似独立同分布每个客户端从每个类别取近似等量样本 split_points np.array_split( indices, num_clients ) else: # 非独立同分布少量客户端垄断部分类别 num_clients_with_class max( 1, int(num_clients * (1 - non_iid_degree)) ) split_points np.array_split( indices, num_clients_with_class ) for client_id, part in enumerate(split_points): client_data_indices[client_id % num_clients].extend(part.tolist()) return [ Subset(dataset, indices) for indices in client_data_indices ]这段代码的核心思路是先按样本标签将所有数据分成类别桶再根据non_iid_degree参数决定每个类别桶如何分配给客户端。当non_iid_degree接近 1 时每个类别只分给少数几个客户端形成典型的“标签偏斜”非独立同分布场景当它接近 0 时每个客户端都能拿到所有类别的样本趋近独立同分布。这个参数直接影响了联邦学习实验的难度——你把non_iid_degree调得越高客户端本地模型偏移就越严重全局聚合后的模型收敛就越慢这可以用来检验不同聚合算法的鲁棒性。需要注意的是np.array_split与np.split的区别前者在不整除时会自动调整每个分片的大小后者要求严格等分否则报错。在客户端数量不能整除类别样本数时必须用array_split。另外client_id % num_clients这行代码保证了客户端数量多于类别数时多出的分片会被轮转分配给前面的客户端不至于浪费样本。3.2 客户端实现本地训练循环与参数序列化客户端的职责非常纯粹接收服务器下发的全局参数用自己的本地数据训练若干轮然后把更新后的参数返回给服务器。理解这个模块要注意一个关键点——客户端和服务器之间传递的不是梯度而是模型权重张量。在横向联邦学习中保护数据隐私的最小单位就是模型参数原始数据永远不离开客户端。# client.py: 联邦客户端核心逻辑 import flwr as fl import torch from collections import OrderedDict def get_parameters(net): 将模型参数转换为有序字典便于传输和聚合 return OrderedDict( {k: v.cpu().numpy() for k, v in net.state_dict().items()} ) def set_parameters(net, parameters): 将接收到的全局参数写回本地模型 params_dict zip(net.state_dict().keys(), parameters) state_dict OrderedDict( {k: torch.tensor(v) for k, v in params_dict} ) net.load_state_dict(state_dict, strictTrue) class FlowerClient(fl.client.NumPyClient): def __init__(self, net, trainloader, testloader): self.net net self.trainloader trainloader self.testloader testloader def get_parameters(self, config): 返回当前模型参数 return get_parameters(self.net) def fit(self, parameters, config): 接收全局参数本地训练返回更新后的参数 set_parameters(self.net, parameters) train_model(self.net, self.trainloader, epochsconfig[local_epochs]) return get_parameters(self.net), len(self.trainloader.dataset), {} def evaluate(self, parameters, config): 在本地测试集上评估模型精度 set_parameters(self.net, parameters) loss, accuracy test_model(self.net, self.testloader) return loss, len(self.testloader.dataset), {accuracy: accuracy}这个类继承自flwr.client.NumPyClient这是 Flower 提供的简化接口它帮我们处理了 NumPy 数组和 PyTorch 张量之间的转换。fit方法是整个联邦学习的核心执行点先调用set_parameters把服务器下发的全局权重覆盖到本地模型再调用train_model在本地数据上训练指定轮次最后把新权重和训练样本数量返回。这里返回的len(self.trainloader.dataset)非常关键因为 FedAvg 算法在聚合时要以样本数量为权重数据量大的客户端在全局模型中的话语权更大。evaluate方法用于在联邦学习过程中周期性评估全局模型的性能。每个客户端用自己的本地测试集评估模型返回精度服务器收集所有客户端的评估结果后取平均。这样做的好处是能观察到全局模型在不同数据分布上的表现差异——如果某个客户端的数据分布偏差很大它的本地评估精度通常会低于其他客户端。get_parameters方法中的v.cpu().numpy()是避免设备不一致的常见做法因为在训练过程中模型可能被移到 GPU 上而传输参数需要先移回 CPU 并转为 NumPy 格式。如果你使用的是 GPU 环境一定要保留这行转换逻辑否则load_state_dict时会报张量设备不匹配的错误。3.3 服务器聚合逻辑FedAvg 的加权平均实现服务器的核心是聚合策略。FedAvg 的全称是 Federated Averaging做法非常简单但极其有效对每个参与客户端上传的参数按照其本地样本数量占总样本数量的比例进行加权平均。这个算法的精妙之处在于它不需要任何额外的通信开销只是简单地对参数做线性组合。# server.py: 自定义聚合策略 import flwr as fl from flwr.server.strategy import FedAvg class WeightedFedAvg(FedAvg): def aggregate_fit(self, server_round, results, failures): 重写聚合方法实现基于样本量的加权平均 results: list of (ClientProxy, FitRes) FitRes包含参数和num_examples if not results: return None, {} # 计算总样本量 total_examples sum( [res.num_examples for _, res in results] ) # 按样本量加权平均 weighted_parameters None for _, res in results: parameters res.parameters weight res.num_examples / total_examples if weighted_parameters is None: weighted_parameters [ (arr * weight) for arr in parameters ] else: weighted_parameters [ acc (arr * weight) for acc, arr in zip(weighted_parameters, parameters) ] # 转回 Parameters 对象 parameters fl.common.parameters_to_ndarrays(result.parameters) for ... # 此处处理参数格式转换 return parameters, {}这段代码把 FedAvg 的数学公式直接翻译成了 Python每个客户端的参数乘以它的样本占比然后累加。num_examples是每个客户端在上传参数时带来的训练样本数服务器依据它计算权重。这种实现方式的优势在于它不依赖任何第三方库纯 Python 列表推导即可完成方便你在此基础上改成 FedProx、FedAdam 等其他聚合算法。aggregate_fit返回的第一个值是聚合后的全局参数第二个是额外的指标字典可以存放聚合过程中计算的元信息。如果你要实现 FedProx只需要增加一个近端项约束要实现带动量更新的聚合则需要在aggregate_fit里维护一个动量缓冲区。这个接口的扩展性设计是 Flower 框架最值得学习的部分——策略类封装了联邦学习的全部逻辑你可以像写普通类一样通过继承和重写来定制算法。3.4 入口脚本如何协调多进程启动入口脚本承担了“导演”的角色读取配置、切分数据、启动服务器、创建多个客户端进程。这里的难点在于进程管理和数据隔离——每个客户端必须只能看到自己那份数据否则就失去了模拟联邦学习的意义。# run_simulation.py: 多进程启动入口 import yaml import multiprocessing from flower.client import start_client from flower.server import start_server def run_client(client_id, data_dir, server_address): 每个客户端进程的入口函数 train_dataset load_partition(data_dir, client_id) net create_cnn() client FlowerClient(net, train_dataset, None) fl.client.start_numpy_client( server_addressserver_address, clientclient, grpc_max_message_length1024 * 1024 * 1024 ) if __name__ __main__: config yaml.safe_load(open(config.yaml)) # 先启动服务器进程 server_process multiprocessing.Process( targetstart_server, args(config,) ) server_process.start() # 再启动多个客户端进程 client_processes [] for client_id in range(config[num_clients]): p multiprocessing.Process( targetrun_client, args(client_id, config[data_dir], config[server_address]) ) p.start() client_processes.append(p) for p in client_processes: p.join()这里有一个值得注意的设计细节grpc_max_message_length被设置成了 1GB。这是一个非常多见的坑——PyTorch 模型的参数转换成 NumPy 后体积可能很大默认的 gRPC 消息大小限制4MB经常导致参数上传失败。很多新手在模拟联邦学习时遇到通信报错排查到最后才发现是这个参数没调大。进程启动顺序也有讲究必须先启动服务器再启动客户端。因为客户端启动后会立即尝试连接服务器完成注册如果服务器还没就绪客户端进程会抛连接异常。在实际调试中可以给server_process.start()后加一个time.sleep(3)让服务器有充分时间初始化端口监听。另外client_processes.join()会让主进程阻塞等待所有客户端退出这保证了服务器进程能收集到所有客户端的训练结果。4. 运行配置与依赖管理从零到训练曲线出现4.1 环境准备与依赖安装这个项目的运行环境要求并不苛刻Python 3.8 以上即可推荐使用虚拟环境隔离依赖避免污染系统 Python。依赖项集中在三块深度学习框架 PyTorch、联邦学习框架 Flower 和配置解析库 PyYAML。如果你有 Nvidia GPU安装对应 CUDA 版本的 PyTorch 能显著加速本地训练如果没有 GPUCPU 版本也能跑通全部代码流程只是每个本地训练轮次会慢一些。# 一键安装依赖 pip install torch torchvision flwr[simulation] pyyaml # 如果需要使用 Ray 后端做并行模拟额外安装 pip install rayflwr[simulation]这个扩展包包含模拟器相关的依赖。如果你不想用多进程而是想用线程模拟可以选择安装ray后端Flower 的模拟模式支持 Ray 来编排多个虚拟客户端但在本地调试时多进程方式更直观且更容易断点调试。安装完成后可以用python -c import flwr; print(flwr.__version__)验证安装是否成功。4.2 配置文件逐项解读config.yaml是控制整个联邦学习实验的参数面板理解每个参数的含义是进行有效实验的前提。# config.yaml num_clients: 10 # 参与联邦学习的客户端总数 fraction_fit: 0.8 # 每轮实际参与训练的客户端比例 batch_size: 32 # 本地训练的批量大小 local_epochs: 3 # 每个客户端本地训练轮数 num_rounds: 20 # 联邦聚合总轮数 learning_rate: 0.01 # 本地训练的学习率 non_iid_degree: 0.7 # 非独立同分布程度 server_address: 127.0.0.1:8080 # 服务器监听地址 dataset: mnist # 数据集名称mnist / cifar10fraction_fit控制每轮参与聚合的客户端比例。在真实场景中由于网络不稳定或设备电量限制并不是所有客户端都能参与每一轮训练。模拟时设置 0.8 意味着 10 个客户端中只有 8 个被随机选中参与本轮这增加了实验的真实性和鲁棒性测试价值。local_epochs是联邦学习和集中式学习最大的区别之一——客户端在本地只训练少量轮次避免在私有数据上过拟合。实践经验是local_epochs设置过大会导致客户端模型漂移不同客户端更新方向不一致全局聚合后模型反而变差。non_iid_degree直接影响实验结论的可信度建议在做算法对比时固定这个参数不变只改变聚合算法变量。而num_rounds控制联邦学习的总通信轮次在单机模拟中一轮通信从启动到完成通常只需要十几秒20 轮足够观察到 MNIST 上从 10% 精度提升到 90% 以上的完整过程。4.3 启动实验并验证训练效果启动实验只需要一条命令python run_simulation.py。服务器进程会打印当前的聚合轮次客户端进程会输出本地训练进度。要验证训练是否正常收敛最直接的方式是观察每轮联邦学习结束后的全局模型精度变化——正常情况是精度逐轮上升曲线先陡后缓。$ python run_simulation.py [server] Round 1/20: 8 clients selected, aggregation done, accuracy0.32 [server] Round 2/20: 7 clients selected, aggregation done, accuracy0.41 [server] Round 3/20: 8 clients selected, aggregation done, accuracy0.52 ... [server] Round 20/20: 8 clients selected, aggregation done, accuracy0.93如果精度在第 3 轮之后不再上升甚至下降优先检查几个地方learning_rate是否过大导致本地模型震荡local_epochs是否过小导致本地模型欠拟合以及non_iid_degree是否过高导致聚合困难。一个实用的调参技巧是先设置non_iid_degree0跑通基线确认代码逻辑正确后再逐步调高非独立同分布程度观察精度变化这样能有效定位问题是出在数据分布还是出在聚合算法上。5. 常见问题与避坑排查单机模拟联邦学习最容易翻车的 5 个瞬间5.1 客户端连接被拒绝进程直接崩溃现象启动run_simulation.py后客户端进程报ConnectionRefusedError或channel connect failed随后进程退出只有服务器进程还在运行。原因客户端启动的速度快于服务器完成端口监听的初始化速度或者服务器进程内部已经异常退出但主进程没有及时发现。解决在启动服务器进程后强制等待若干秒确保监听端口已就绪。如果问题仍然存在用netstat -an | grep 8080检查端口是否被占用。另外服务器地址必须使用127.0.0.1:8080而不是localhost:8080因为localhost在某些操作系统上会解析为 IPv6 地址::1导致 IPv4 连接失败。我一般在调试阶段会把server_process.start()后的等待时间设为 5 秒虽然不够优雅但足够稳定。5.2 聚合时参数全变为零模型精度恒为随机水平现象联邦学习正常启动但全局模型精度始终在 10% 附近MNIST 10 分类任务随机水平打印服务器日志发现聚合后的参数数组全部为 0。原因这是 FedAvg 实现中最经典的翻车点。在加权平均时忘记将 NumPy 数组乘以权重而是直接赋值导致每次聚合结果来自最后到达的客户端参数。还有一种可能是初始化时把参数设置成了全零但模型结构包含偏置项时全零初始化会让梯度更新失效。解决在aggregate_fit中打印第一个客户端的参数形状和值范围确认参数非零。检查加权累加逻辑确保weighted_parameters是逐元素累加而不是重新赋值。核心代码如下if weighted_parameters is None: weighted_parameters [arr * weight for arr in params] else: weighted_parameters [acc arr * weight for acc, arr in zip(weighted_parameters, params)]5.3 训练完第一轮后客户端报 Tensor 设备不匹配错误现象本地训练正常完成但当客户端调用get_parameters准备上传参数时报RuntimeError: Expected all tensors to be on the same device, but found at least two devices。原因模型在 GPU 上训练后部分参数仍留在 GPU而state_dict().items()迭代出来的张量混有 CPU 和 GPU。这在多 GPU 或 GPU 空闲时切换到 CPU 的场景中尤其常见。解决在返回参数前强制调用v.cpu().numpy()确保所有参数张量被移回 CPU。这也是为什么前面章节中get_parameters方法特别强调cpu()转换。如果项目中使用了torch.nn.DataParallel包装模型状态字典的键会带有module.前缀需要在set_parameters时做前缀兼容。5.4 客户端本地训练很快但通信时间非常长现象日志显示每个客户端训练只花了 2 秒但整个联邦轮次完成却花了 30 秒大量时间消耗在参数上传和下发阶段。原因默认的 gRPC 消息限制导致参数在传输前被压缩到很小但反序列化开销大另一种常见情况是模型很大但grpc_max_message_length设置过小参数被分片传输后频繁触发重传机制。解决在客户端和服务器的启动参数中同时设置grpc_max_message_length1024*1024*1024给参数传输留出充足空间。此外减少fraction_fit参数可以降低每轮通信的客户端数量在调试算法阶段能显著加速实验迭代速度。我一般先用 5 个客户端跑通流程确认算法正确后再扩展到 10 个以上。5.5 联邦学习精度低于集中式训练但找不出原因现象相同的模型结构、批次大小、轮次设置联邦学习最终精度比集中式训练低 2% 到 5%且调整学习率没有明显改善。原因这是联邦学习本身的特性而非代码 Bug。客户端本地训练轮次过多会导致模型参数偏离全局最优方向非独立同分布数据会放大这种偏差。此外每轮只有部分客户端参与聚合相当于每个通信轮次只消费了部分数据模型收敛所需的有效信息量被稀释。解决先降低local_epochs到 1看精度是否提升然后把fraction_fit提高到 1.0让所有客户端每轮都参与排除采样偏差。如果两步都不奏效考虑在本地训练中加入 L2 正则化或使用 FedProx 的近端项约束。需要注意的是联邦学习的精度略低于集中式训练是正常现象如果差距在合理范围内说明实现正确。6. 二次开发上手从 FedAvg 扩展到 FedProx、个性化联邦学习与多机部署拿到这份源码后二次开发的方向取决于你想解决什么问题。最直接的是更换聚合算法——在server.py中继承WeightedFedAvg并重写aggregate_fit方法加入 FedProx 的近端项惩罚只需要在客户端本地训练时增加一个正则项代码改动不超过 10 行。FedProx 的动机是限制本地模型与全局模型的偏差在train_model函数里加入以下逻辑# 在本地训练的损失函数中加入近端项 def proximal_loss(model, global_params, mu0.01): proximal_term 0.0 for local_param, global_param in zip( model.parameters(), global_params ): proximal_term ((local_param - global_param) ** 2).sum() return mu * proximal_term / 2把proximal_loss加到原有损失函数上mu控制约束强度。mu太大模型更新缓慢太小起不到约束作用实验时从 0.01 开始调。这个改动让联邦学习在非独立同分布数据上的收敛稳定性明显提升是面试时能讲出技术深度的关键话题。第二个方向是从横向联邦过渡到个性化联邦学习。横向联邦学习假设所有客户端共享一个全局模型但现实中各客户端数据分布差异大时单个全局模型难以同时适配所有客户端。可以借鉴 Per-FedAvg 的思路让每个客户端在收到全局模型后在自己的本地数据上多微调几步产生个性化模型。在FlowerClient.fit中修改返回值# 个性化联邦学习返回前增加本地微调 def fit(self, parameters, config): set_parameters(self.net, parameters) # 在返回值之前做本地适配 for _ in range(config[personalization_epochs]): train_one_epoch(self.net, self.personal_trainloader) train_model(self.net, self.trainloader, epochsconfig[local_epochs]) return get_parameters(self.net), len(self.trainloader.dataset), {}第三个方向是验证算法后往多机部署迁移。Flower 的接口设计保证了这一点本地模拟时用多进程启动多个客户端迁移到多机时只需要把每个客户端进程部署到不同机器上客户端中start_numpy_client的server_address改为服务器的实际 IP 和端口服务器代码完全不需要改动。这个过程让我深刻体会到好的抽象带来的好处——业务代码与通信细节完全解耦模拟环境与生产环境共享同一套核心逻辑。从那以后我每次做联邦学习实验都强制先在本地用多进程模拟跑通全流程确认聚合算法正确后再上真实集群这套习惯帮我避免了大量在分布式环境下无从排查的玄学 bug。希望这个模拟项目的代码结构和踩坑记录能帮你在联邦学习的路上省下几天的调试时间。本文还有配套的精品资源点击获取
