使用NVIDIA FLARE构建并比较非独立同分布CIFAR-10上的FedAvg与FedProx联邦学习分步指南
本教程详细介绍了如何使用NVIDIA FLARE构建高级联邦学习实验，在非独立同分布CIFAR-10数据集上比较FedAvg和FedProx算法。通过狄利克雷分布模拟客户端的标签不平衡，并使用NVFlare Job API定义和启动联邦任务，Client API处理本地训练和模型交换。文章提供了完整的代码实现和实验结果可视化。
文章情报
要点
- 使用NVIDIA FLARE构建联邦学习实验，比较FedAvg与FedProx。
- 采用狄利克雷分布（alpha=0.3）将CIFAR-10数据划分为3个非独立同分布客户端。
- 通过NVFlare Job API配置服务器端协调，Client API实现客户端本地训练与模型上传。
- 实验结果表明FedProx在非独立同分布数据上可能优于FedAvg，通过全局测试准确率曲线对比。
为什么重要
这条新闻值得关注，因为使用NVIDIA FLARE构建联邦学习实验，比较FedAvg与FedProx。
技术影响
可能影响模型选型、推理成本、产品能力和评测基准。
在联邦学习领域，处理非独立同分布（non-IID）数据是一个核心挑战。本文提供了一个完整的实践教程，基于NVIDIA FLARE框架，在CIFAR-10数据集上对比了两种经典联邦学习算法：FedAvg和FedProx。教程面向有一定深度学习基础的开发者，通过实际的代码示例展示了如何构建和运行联邦学习实验。
首先，我们准备实验环境。教程要求安装nvflare、torch、torchvision和matplotlib等库。我们需要定义实验参数，包括客户端数量（3个）、通信轮次（5轮）、本地训练轮次（1轮）、批大小（64）、学习率（0.01）以及控制非独立同分布程度的狄利克雷参数alpha（0.3）。数据集CIFAR-10被下载到指定目录，供所有模拟客户端共享。
为了模拟现实场景中的标签分布不均，教程使用狄利克雷分布将训练数据划分给各个客户端。通过设置固定的随机种子，每个客户端进程可以独立地生成相同的全局划分。这样，每个客户端得到的数据标签比例各不相同，更接近真实联邦学习环境。
客户端训练脚本是实验的核心部分。我们定义了一个小型CNN模型（包含两个卷积层和两个全连接层），用于CIFAR-10的10分类任务。脚本使用NVFlare的Client API进行初始化，接收服务器下发的全局模型参数，并在本地数据上进行训练。对于FedProx算法，我们引入了近端项（proximal term），通过在损失函数中添加模型参数与全局模型参数的L2距离惩罚项来控制局部更新偏离全局模型的程度。教程中的mu参数设置为0.0（FedAvg）和0.1（FedProx）。
在服务器端，我们使用NVFlare的Job API创建FedAvgJob。通过ScriptRunner将客户端脚本附加到每个模拟站点（site-1、site-2、site-3），并传递统一的实验参数。然后，我们分别运行两个实验：标准FedAvg和带近端项的FedProx。每个实验结束后，训练日志以CSV格式保存，包含每轮通信后的全局测试准确率。
最后，我们加载两个实验的准确率记录，绘制学习曲线。曲线展示了随着通信轮次的增加，全局模型在非独立同分布CIFAR-10测试集上的表现。通过对比，我们可以直观地观察到FedProx相对于FedAvg在标签不平衡场景下的潜在优势。此外，教程还展示了如何定位最终的全局模型检查点，以便后续分析或复用。
总之，本教程提供了一个端到端的联邦学习工作流程，涵盖了从数据划分、客户端训练、服务器协调到结果可视化的完整步骤。读者可以基于此模板进一步探索其他联邦学习算法或调整参数，以应对更复杂的非独立同分布数据场景。