FedLab
基于PyTorch的联邦学习研究框架,支持单机到多机多进程4种通信模式,内置17+种联邦学习算法和F
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
基于PyTorch的联邦学习研究框架,支持单机到多机多进程4种通信模式,内置17+种联邦学习算法和F
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
想象这样一个场景:三家三甲医院的影像科医生都想训练一个肺癌早筛 AI 模型,但患者数据涉及隐私,任何一家医院都无法直接把 CT 片发给其他医院共享。传统的集中式训练在这里碰壁——不是技术不够,而是法规和隐私不允许。
联邦学习(Federated Learning) 就是为解决这个问题而生。它的核心思想是"数据不动,模型动":各医院在本地用自家数据训练模型,只将模型参数(而非原始数据)上传到中央服务器,由服务器聚合后分发新参数。如此循环迭代,最终所有医院共同享有一个性能强大的模型,同时患者数据始终留在本地。
FedLab 正是这样一个专为研究者和开发者设计的联邦学习算法实验框架,基于 PyTorch 实现,提供了从单机模拟到多进程/多机通信的全套抽象,让联邦学习算法的研发周期大幅缩短。
FedLab 由鹏城实验室(Peng Cheng Laboratory)下属的 SMILELab 团队开发和维护。鹏城实验室是国内人工智能领域的顶级研究机构,其在分布式计算、隐私计算方向有深厚积累。FedLab 项目起始于 2021 年,目前已发布 v1.3.0 版本,GitHub 累计获得 825 颗星、142 次 Fork,说明其在学术界和工业界都有一定影响力。
项目的核心维护者包括曾敦(Dun Zeng)、梁思齐(Siqi Liang)和胡向静(Xiangjing Hu)三位研究者,他们在联邦学习领域发表了多篇相关论文,FedLab 本身也可以视为这些研究成果的工程化实现。
FedLab 的设计亮点在于将联邦学习系统的通信架构抽象为四个层次,从简单到复杂,逐一满足不同研究场景的需求:
1. Standalone(单机模式):在同一进程中模拟多个客户端,所有参与方共享内存通信,开箱即用,最适合快速验证算法想法。MNIST 图像分类示例只需几十行代码即可运行。
2. Serial Trainers(串行训练器):在同一台机器上通过多进程模拟多个客户端进程,适合研究进程间通信(IPC)场景,支持比单机模式更真实的资源竞争模拟。
3. Asynchronous FL(异步联邦学习):客户端不等彼此,按各自节奏上传参数,服务器实时聚合。这是 FedLab 的特色模式之一,相比同步联邦学习,异步模式在客户端计算能力差异大或网络不稳定时表现更好。
4. Hierarchical Hybrid(分层混合):引入边缘服务器层,形成"客户端→边缘服务器→中央服务器"两层架构。这种设计模拟了真实企业场景:各分公司是客户端,省级服务器是边缘节点,总部是中央服务器。模型在边缘层做初步聚合,大幅降低中央服务器通信压力。
框架的核心类分为两大族:Client(客户端) 和 Server(服务器)。客户端侧负责本地训练和数据分区,服务器侧负责参数聚合和全局调度。network_manager.py 和 communicator.py 封装了底层的 socket 通信逻辑,上层算法开发者无需关心网络细节。
FedLab 不仅仅是一个框架,更是一个算法百宝箱。在 fedlab/contrib/algorithm/ 目录下,预置了 17 种主流联邦学习算法:
| 算法 | 类别 | 核心特点 |
|---|---|---|
| FedAvg | 同步联邦 | 基准算法,参数加权平均 |
| FedProx | 同步联邦 | 加入正则项处理客户端异构性 |
| FedNova | 同步联邦 | 校正本地训练步数的影响 |
| FedMGDA+ | 同步联邦 | 基于梯度下降的多目标优化 |
| FedOpt | 同步联邦 | 服务器端引入优化器(如 Adam) |
| FedAvgM | 同步联邦 | 引入服务器动量加速收敛 |
| Ditto | 个性化联邦 | 每个客户端保留个性化模型 |
| SCAFFOLD | 同步联邦 | 引入控制变量纠正梯度偏移 |
| FedDyn | 同步联邦 | 动态正则化实现收敛保证 |
| FedAsync | 异步联邦 | 异步聚合策略 |
| FedBuff | 异步联邦 | 缓冲区机制控制异步延迟 |
| CFL | 异步联邦 | 基于聚类的客户端分组 |
| IFCA | 聚类联邦 | 迭代式聚类联邦,平均同时聚类 |
| PowerofChoice | 客户端选择 | 选择高梯度贡献的客户端参与轮次 |
| qFedAvg | 公平性联邦 | 考虑客户端数据量差异的加权 |
每种算法都有完整的训练流程封装,可以直接继承 BasicClient 和 BasicServer 类,替换聚合逻辑即可。代码结构遵循 PyTorch 的模块化习惯,理解和二次开发成本较低。
联邦学习的核心瓶颈之一是通信开销。当模型参数量达到百万甚至十亿级别时,每轮传输的梯度数据量非常可观。FedLab 在 fedlab/contrib/compressor/ 中实现了完整的通信压缩工具链:
这些压缩工具被设计为可插拔的中间件,集成到通信管道中,对上层算法完全透明。
研究联邦学习时,最大的痛点之一是"训练过程不透明"——模型在多个节点上训练,到底收敛了没有?各客户端的损失曲线差异大吗?
FedLab 提供了内置的可视化组件 FedBoard(位于 fedlab/board/),基于 Python-Flask + Vue.js 构建,提供了一个轻量级的 Web Dashboard,可以实时监控:
FedBoard 是可选组件,可以通过 pip 单独安装,适合在调试算法时使用,生产环境中通常会替换为专业的实验监控系统。
FedLab 的技术选型非常清晰:
代码结构采用标准的 Python 包布局,fedlab/core/ 放置核心抽象类,fedlab/contrib/ 放置扩展算法,fedlab/models/ 放置预置模型(CNN、MLP、RNN),fedlab/utils/ 放置数据分区、序列化等工具函数。
代码质量方面,项目有完整的单元测试套件(tests/ 目录),CI/CD 通过 GitHub Actions,覆盖率通过 codecov 监控。文档质量较高(docs/ 目录 + 官方 readthedocs 站点),提供 4 个 Jupyter Notebook 教程,覆盖通信、定制化、流水线等核心主题。
部署 FedLab 有两种路径:
路径一(推荐):pip 安装
pip install fedlab
官方已发布至 PyPI,适合快速试用。但需注意依赖 PyTorch 1.7.1+,且默认不含 GPU 优化版。
路径二:Docker 容器
docker/Dockerfile 提供了基于官方 PyTorch 镜像的构建示例,需要手动指定 CUDA 版本和 PyTorch 版本。镜像构建时从清华 conda 镜像源安装依赖,但文档说明较为简略,非 Docker 熟练用户可能需要花费额外时间排错。
路径三:从源码安装
git clone https://github.com/SMILELab-FL/FedLab
cd FedLab && pip install -e .
适合需要修改框架源码的研究者。
硬件方面,CPU 可运行 Standalone 模式的基本示例,但涉及真实模型训练时建议使用 NVIDIA GPU(CUDA 10.1+)。内存建议 8GB 以上,磁盘 5GB(包含数据集)。
FedLab 并非生产级联邦学习系统,以下局限需要关注:
非安全通信:当前实现未集成差分隐私或安全多方计算(MPC),数据隐私保护依赖协议层面的信任假设,不适合直接用于有严格合规要求的生产环境。
Docker 体验不完善:docker-compose.yml 不存在,K8s 支持为零,容器化部署更多是"能用"而非"好用"。
文档语言:主要文档为英文,对中文用户存在一定门槛。不过由于代码注释清晰,结合 Jupyter Notebook 教程,上手难度可控。
版本兼容:要求 Python >= 3.6 和 PyTorch >= 1.7.1,在更新的 PyTorch 2.x 环境下部分功能可能存在兼容性问题,建议使用 1.x 版本。
FedLab 的价值在于降低了联邦学习研究的入门门槛。在它出现之前,研究者要验证一个新算法,往往需要从头搭建通信框架、处理数据分区、实现聚合逻辑,代码量动辄数千行。FedLab 将这些共性工作抽象为可复用的组件,研究者可以专注于算法逻辑本身。
从生态角度看,FedLab 与 PySyft(OpenMined 出品的隐私计算框架)、Flower(通用的联邦学习平台)形成了互补:PySyft 侧重隐私保护机制,Flower 侧重生产部署,FedLab 则侧重算法研究和教学实验。三者覆盖了联邦学习从研究到落地的不同阶段。
截至目前,FedLab 仍在活跃维护中(2024 年仍有提交),项目的发展方向包括更好的异构数据支持、更丰富的个性化联邦学习算法,以及与主流 ML 平台(如 PyTorch Lightning)的集成。对于想在联邦学习方向开展研究或原型的团队,FedLab 是值得优先考虑的技术选型。
关键信息一览