TensorFlow1.15中tf.nn.softmax输出nan/inf，裁剪logits到-15~15可解决的原因

1. 核心诱因：TensorFlow 1.15 GPU版本的Softmax CUDA内核存在已知数值稳定性缺陷 #

你在交互式环境测试时调用的是CPU版本Softmax实现，该实现完整搭载了标准Softmax数值稳定优化逻辑：计算前先减去当前样本logits的最大值，避免指数运算溢出，因此即使输入极值也能输出正常结果。而你实际训练时调用的是GPU版CUDA Softmax内核，该内核在TF 1.15版本中存在边界处理漏洞，当单样本的logits极差（最大值减最小值）超过约30时，会触发未定义的数值计算错误，生成nan/inf。

2. 裁剪阈值的合理性 #

• 将logits裁剪到[-15, 15]区间时，单样本logits的最大极差为30，刚好落在GPU Softmax内核的安全计算范围内，因此不会触发错误
• 当裁剪范围放宽到[-20, 20]时，最大极差达到40，超过了安全阈值，因此依然会报错

额外优化建议 #

如果不想裁剪logits损失精度，可以替换tf.nn.softmax + 后续交叉熵的组合为tf.nn.softmax_cross_entropy_with_logits_v2，这个API的GPU实现做了完善的数值稳定处理，不会出现类似溢出问题，同时还能减少计算开销。

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