TensorFlow GPU训练损失与准确率无变化问题排查求助

核心原因 #

• 你使用的tf.keras.layers.experimental.preprocessing.StringLookup是TensorFlow的早期实验性实现，在disable_eager_execution开启的静态图模式+GPU运行环境下，存在梯度传递中断的bug：反向传播的梯度无法穿过这一层传递到后续的Embedding层，导致Embedding权重全程没有更新，自然损失和准确率不会变化。而CPU环境下该层的静态图算子实现不同，梯度可以正常传递，所以CPU训练正常。
• SageMaker官方TensorFlow镜像默认开启FP16混合精度训练，你的模型结构中经过L2归一化的Dot层输出范围在[-1,1]，后续接Sigmoid激活时很容易进入饱和区，GPU混合精度下会出现梯度下溢为0的问题，进一步加剧了权重不更新的现象。
• 静态图模式下GPU侧的权重更新默认不会主动同步到CPU内存，你在回调中通过model.layers[2].get_weights()[0]获取的始终是初始权重，也会导致你观测到的相似度结果没有变化。

修复方案 #

• 直接删除disable_eager_execution相关代码，使用TensorFlow 2.x默认的即时执行模式跑GPU训练，你CPU上移除该段代码后运行正常，GPU上移除后可直接规避静态图的兼容性问题，当前版本TF的eager模式GPU训练速度已经和静态图基本持平。
• 把字符串转索引的逻辑从模型层移到数据预处理阶段：提前将训练数据里的company_id、payee字段映射为对应的int索引后存储，训练时直接读取int值输入模型，彻底规避StringLookup层的兼容问题。
• 关闭混合精度训练，或者给Dot层去掉normalize=True参数，手动实现L2归一化时加入1e-8的epsilon参数避免除零，缓解梯度消失问题。
• 恢复训练数据的shuffle操作，避免模型拟合样本顺序特征。

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