本地GTX1650训练CIFAR-10模型loss为NaN，Colab正常问题排查

1. 硬件差异导致的混合精度兼容性问题 #

GTX 1650属于Turing架构但没有Tensor Cores，而Colab常用的Tesla T4/V100/A100等GPU配备了Tensor Cores——TensorFlow的mixed_float16策略在有Tensor Core的硬件上会做专门的数值优化，无Tensor Core的GPU则依赖软件模拟，数值稳定性差异极大。

• 解决方法：
  强制本地使用mixed_float32或直接关闭混合精度，验证是否还会出现NaN：

  # 替换原混合精度设置
  mixed_precision.set_global_policy('mixed_float32')
  # 或者完全关闭混合精度
  # mixed_precision.set_global_policy('float32')
  

2. 损失函数的数值精度问题 #

float16的数值范围远小于float32，若损失函数的计算逻辑未适配无Tensor Core的硬件，容易出现溢出或下溢导致NaN。比如分类任务中，logits范围过大时，SparseCategoricalCrossentropy在float16下极易出问题。

• 解决方法：
  强制损失函数以float32精度计算：

  loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, dtype='float32')
  

  或者自定义损失函数手动转换精度：

  def custom_loss(y_true, y_pred):
      y_pred = tf.cast(y_pred, tf.float32)
      return tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)
  

3. 小批量BatchNorm的数值稳定性问题 #

本地GTX 1650显存通常比Colab GPU小，你大概率用了更小的batch size。小批量下BatchNorm的统计量（均值、方差）本身就不稳定，再加上float16的精度限制，容易出现方差趋近于0的情况，直接导致NaN。

• 解决方法：
  • 若显存允许，适当增大batch size；
  • 用LayerNorm替代BatchNorm，减少对批量统计量的依赖；
  • 手动初始化BatchNorm的gamma参数为小正值，避免初始数值异常：

    tf.keras.layers.BatchNormalization(gamma_initializer=tf.keras.initializers.Constant(0.1))
    

4. LossScaleOptimizer的动态缩放失效 #

虽然你用了LossScaleOptimizer，但无Tensor Core的GPU上，动态缩放的自适应逻辑可能出现偏差，导致缩放因子过大或过小，引发数值溢出/下溢。

• 解决方法：
  改用静态损失缩放，设置固定的缩放因子：

  base_optimizer = tf.keras.optimizers.Adam(clipnorm=1.)
  optimizer = tf.keras.mixed_precision.LossScaleOptimizer(base_optimizer, dynamic=False, initial_scale=65536)
  

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