SecretFlow项目中FedProx算法训练BN层模型的问题分析与解决方案

问题背景

在SecretFlow项目的联邦学习实现中，FedProx算法在训练带有批量归一化(Batch Normalization, BN)层的神经网络模型时出现了不收敛的问题。经过开发者排查，发现这是由于算法实现中缺少设置模型训练模式的语句导致的。

技术细节分析

批量归一化是深度学习中常用的技术，它在训练和推理阶段有不同的行为模式：

1. 训练模式：BN层会计算当前批次的均值和方差，并使用它们来归一化数据，同时更新运行统计量
2. 评估模式：BN层使用训练期间积累的运行统计量，而不是当前批次的统计量

在PyTorch框架中，需要通过model.train()和model.eval()来显式切换这两种模式。而SecretFlow的FedProx实现中遗漏了这一关键设置，导致BN层始终处于评估模式，无法正确更新统计量，最终造成模型无法收敛。

问题复现与验证

开发者使用CIFAR-100数据集进行了验证实验，数据分布在10个参与方之间，采用狄利克雷分布(β=0.5)进行划分。网络结构包含：

1. 特征提取部分：多层卷积+BN层+ReLU的组合
2. 分类部分：2层全连接网络

实验参数配置为：

• 优化器：Adam (学习率5e-4)
• 训练轮次：75
• 权重衰减：1e-4
• FedProx参数μ：0.001

实验结果表明，添加self.model.train()语句后，模型能够正常收敛，验证了问题的根源。

解决方案

修复方案非常简单但关键：在FedProx的train_step方法开始时添加模型训练模式的设置：

def train_step(self, weights: np.ndarray, cur_steps: int, train_steps: int, **kwargs):
    assert self.model is not None, "Model cannot be none, please give model define"
    self.model.train()  # 关键修复
    ...

经验总结

这个案例提醒我们：

1. 在实现联邦学习算法时，需要考虑底层框架的特定行为
2. 对于包含BN层的模型，训练模式设置是不可忽视的细节
3. 分布式训练中的问题有时表现与集中式训练不同，需要特别关注
4. 模型聚合时需要注意处理BN层特有的参数(如num_batches_tracked)

SecretFlow团队已经确认将在后续版本中修复这一问题，这对于使用FedProx算法训练复杂神经网络模型的用户来说是一个重要的改进。