TF 2.4.1 Keras分布式早停训练的指标重复计算与图异常问题

基础环境与已知现象 #

• 训练运行环境：Google AI Platform，基于TensorFlow 2.4.1版本的Keras框架开发训练脚本
• 初始集群配置：1台master节点 + 1台worker节点，所有节点规格为complex_model_m_gpu，单节点搭载4块NVIDIA Tesla K80 GPU
• 预设逻辑：基于每个epoch计算得到的recall@k指标，配置早停触发规则
• 运行观测到的异常：
  • 单个epoch周期内会输出两次recall@k计算值，指标提升判断分成两个完全独立的并行流程运行，两个流程之间没有任何信息交互：例如epoch1分别产出0.13306、0.12903两个指标值，到epoch3时，两套判断逻辑独立和各自留存的上一轮基准值比对——0.17与0.13306对比、0.11与0.12903对比，没有交叉同步
  • TensorBoard中的标量指标图展示混乱
• 对照验证结果：仅保留1台master节点、移除worker节点后运行相同训练任务，单个epoch仅产出1个指标值，TensorBoard展示恢复正常；初步判断异常和多节点下的分布式策略配置错误有关

异常根因解释 #

• 双指标值、早停双轨道判断的核心原因：当前训练脚本配置的tf.distribute.MirroredStrategy仅面向单节点多GPU场景设计，无法适配多节点的集群拓扑。在2节点的集群环境下，两个节点会各自启动一套完全独立、没有跨节点通信的Keras训练循环——每个节点独立加载全量验证集、独立计算自身得到的recall@k值、独立运行早停回调的判断逻辑，最终就形成了两套互不感知的指标计算、指标判断流程。
• TensorBoard展示混乱的原因：两个节点的训练进程都会在训练过程中，按照自己的epoch进度往日志文件写入相同tag的recall@k标量值，日志中同一tag下交错存在两个节点产出的数值，TensorBoard读取日志时将所有同tag数值按写入时间顺序连线，自然会出现图线跳变、展示混乱的问题。

修复方案 #

• 替换适配多节点的分布式策略：将原有单节点多卡使用的tf.distribute.MirroredStrategy，替换为支持多节点集群同步的tf.distribute.MultiWorkerMirroredStrategy。这套策略会自动完成多节点的训练流程全局协调：所有节点的验证集被自动拆分，每个节点仅计算自身数据分片上的指标中间值，最终由全局协调流程将所有节点的中间值聚合为覆盖全量验证集的全局指标值，从根源上避免每个节点独立计算全量指标的问题。
• 调整回调逻辑：替换多节点策略后，Keras内置的早停回调、TensorBoard回调只会接收全局聚合后的唯一指标值，早停判断会基于全局指标做统一的提升校验，不会再出现两套独立判断轨道的问题；TensorBoard写入日志的也只有统一的全局指标值，标量图展示会恢复正常。

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