TFF联邦微调预训练模型时多轮训练后测试准确率恒定无提升问题求助

问题产生原因 #

• 模型输出层配置错误：当前输出层仅使用Dense(3)未添加对应激活函数，也没有在损失函数中配置from_logits=True，模型输出的类别概率大小关系未发生变化，仅置信度调整，因此损失持续下降但分类预测结果不变，准确率恒定。
• TFF模型结构不匹配：构建iterative_process时传入的model_fn返回的模型结构与你自定义的create_keras_model结构不一致，导致预训练权重、可训练层配置没有正确赋值到全局训练状态中，分类头实际没有参与训练。
• 联邦训练超参不合理：客户端本地训练轮次过少、学习率过低，分类头的更新幅度太小，不足以改变样本的预测类别，导致准确率没有变化。
• 数据分布问题：联邦客户端数据为极端非IID分布，每个客户端仅包含单类或少数类别数据，本地训练后聚合的全局模型无法学习到有效的多分类边界。

调整实现方案 #

1. 修正模型结构与损失配置 #

在输出层添加对应激活函数，保证和损失函数匹配，示例修改如下：

def create_keras_model():
    baseModel = tf.keras.models.load_model(r"path\to\model")
    headModel = baseModel.output
    # 多分类任务添加softmax激活，和categorical_crossentropy损失对应
    model_output = tf.keras.layers.Dense(3, activation='softmax')(headModel)
    model = tf.keras.Model(inputs=baseModel.input, outputs=model_output)
    for layer in baseModel.layers:
        layer.trainable = False
    return model

如果使用from_logits模式，则保留输出层无激活，损失函数配置为tf.keras.losses.CategoricalCrossentropy(from_logits=True)即可。

2. 对齐TFF迭代器的模型结构 #

保证构建iterative_process使用的model_fn返回的模型结构、可训练层配置和create_keras_model完全一致，避免权重赋值错位导致的训练无效。

3. 调整联邦训练超参 #

• 提高客户端本地训练epoch数到3~5，保证分类头有足够的更新量
• 将客户端优化器的学习率调整到1e-3~1e-2区间，避免更新步长过小
• 待分类头训练收敛后，可以解冻基模型的顶部2~3层参与训练，进一步提升模型拟合效果

4. 先做集中式训练验证 #

暂时去掉联邦逻辑，直接用集中式的方法训练你构建的Keras模型，如果集中式训练下准确率仍然不上升，说明是模型或数据预处理的问题，先排除框架层面的影响。

5. 优化非IID数据训练 #

如果数据为极端非IID分布，可调整客户端采样策略，保证每轮参与训练的客户端覆盖所有3个类别，也可在本地损失中添加L2正则项，避免客户端过拟合本地数据。

内容的提问来源于stack exchange，提问作者seni