使用 Hugging Face 与 Flower 的联邦学习

Hugging Face 提供了一份技术指南,介绍如何将 Flower 框架集成到 Transformer 模型中以实现联邦学习(FL)。此集成使得在多个客户端上对预训练模型进行微调成为可能,而无需共享原始数据,从而提升隐私和数据安全性。

使用 Flower 的联邦学习架构

联邦学习使得在多个去中心化客户端和中心服务器之间训练全局模型成为可能。各客户端不在单一位置聚合原始数据,而是在本地使用自己的数据训练模型,并仅将模型参数传回服务器。服务器随后使用预定义的策略聚合这些参数,以更新全局模型。

在提供的实现中,流程遵循以下步骤:

  1. 本地训练:客户端使用自己的数据集进行本地训练。
  2. 参数交换:客户端通过 get_parameters 方法将更新后的参数发送给服务器。
  3. 全局聚合:服务器使用如 FedAvg(联邦平均)等策略聚合所有客户端的参数,将全局权重定义为每轮所有客户端权重的平均值。
  4. 模型分发:服务器通过 set_parameters 方法将更新后的全局参数发送回客户端。

技术实现细节

模型与数据集

实现使用 distilBERTdistilbert-base-uncased)作为基础模型,通过 Hugging Face 的 AutoModelForSequenceClassification 加载,用于二分类序列任务。目标任务是对 IMDB 数据集 进行情感分析,模型被训练以判断电影评分是正面还是负面。

Hugging Face 工作流

标准的 Hugging Face 流程用于数据准备和训练:

  • 数据处理:使用 datasets 库获取 IMDB 数据集,然后使用 AutoTokenizer 进行分词,并加载到 PyTorch DataLoader 对象中。
  • 训练循环:使用 AdamW 优化器实现标准的 PyTorch 训练循环。
  • 评估:在测试阶段使用 evaluate 库计算准确率和损失指标。

Flower 客户端 (IMDBClient)

为了将 Hugging Face 模型与 Flower 框架连接,创建了一个继承自 flwr.client.NumPyClient 的自定义客户端类。该类实现了四个关键方法:

  • get_parameters:提取模型参数为 NumPy 数组,以便传输到服务器。
  • set_parameters:使用从服务器接收的参数更新本地模型的状态字典。
  • fit:执行本地训练函数(train),并返回更新后的参数以及使用的样本数量。
  • evaluate:运行本地测试函数(test),并返回损失和准确率指标。

服务器配置与聚合

为了协调联邦过程,初始化了一个带有特定聚合策略的 Flower 服务器。实现使用 fl.server.strategy.FedAvg,并将 fraction_fit=1.0fraction_evaluate=1.0 配置为所有客户端在每轮训练和评估中均参与。

为处理分布式指标,实现了一个 weighted_average 函数。该函数通过根据每个客户端贡献的样本数量对其指标加权,计算全局准确率和损失,从而确保全局性能度量具有代表性。

框架兼容性

虽然示例使用了 PyTorch,但指南指出相同的联邦学习工作流也可以使用 TensorFlow 实现。Flower 的仿真功能(flwr['simulation'])同样可用于在单一环境(如 Google Colab)中模拟联邦环境,以进行测试。

Sources