CNN模型在TensorFlow CPU与GPU上的预测差异问题咨询

这种跨设备推理不一致的坑我之前踩过好几次，大概率是浮点精度、算子实现差异或者数值溢出导致的，咱们一步步来排查解决：

1. 优先检查浮点精度设置 #

GPU训练时很多框架默认会开启**混合精度（FP16/FP32混合）**来加速，但CPU通常只支持FP32计算。如果训练时用了混合精度，导出的模型可能携带FP16的参数或计算逻辑，CPU处理时就容易出现数值溢出变成NaN。

解决办法：

• 导出模型前，强制将模型参数转换为FP32：

  import tensorflow as tf
  # 加载训练好的模型
  model = tf.keras.models.load_model('your_fcn_model.h5')
  # 保存为FP32精度的模型
  model.save('fcn_model_fp32.h5')
  # 再导出为pb格式
  tf.saved_model.save(model, 'fcn_pb_model')
  

• CPU推理前，设置全局浮点精度为FP32：

  tf.keras.backend.set_floatx('float32')
  

2. 排查算子的GPU/CPU实现差异 #

有些算子在GPU和CPU上的底层实现逻辑并不完全一致，比如：

• 某些自定义激活函数（比如早期版本的Swish）
• BatchNormalization的moving均值/方差更新逻辑
• 池化、上采样层的边缘处理

解决办法：

• 检查模型结构，替换掉可能有兼容性问题的算子：比如把自定义Swish换成TF官方的tf.keras.activations.swish
• 若用了BatchNorm，确保导出模型时已经保存了训练好的moving_mean和moving_variance（优先用tf.saved_model.save而非手动导出pb）
• 在CPU环境下对模型进行1-2个epoch的微调（用少量训练数据），让模型适配CPU的算子计算逻辑

3. 处理数值溢出/下溢 #

GPU的浮点计算动态范围比CPU更宽，训练时某些中间值在GPU上是合法的，但CPU计算时会溢出成NaN（比如极大的权重乘以输入导致数值爆炸）。

解决办法：

• 检查模型权重的分布：打印各层权重的最大值/最小值，如果有极端值（比如大于1e4或小于-1e4），考虑添加梯度裁剪（训练阶段）：

  optimizer = tf.keras.optimizers.Adam(clipvalue=1.0)
  model.compile(optimizer=optimizer, loss='your_loss_function')
  

• 推理时对输入和中间层输出做clip限制：

  input_data = tf.clip_by_value(input_data, -10.0, 10.0)
  predictions = model(input_data)
  # 或者对模型输出做clip
  predictions = tf.clip_by_value(predictions, 0.0, 1.0)
  

4. 核对CPU环境的依赖版本 #

如果CPU环境的TensorFlow（或其他框架）版本和GPU训练时不一致，可能会出现算子兼容性bug。比如某些TF2.x版本的CPU算子对特定层的处理有问题。

解决办法：

• 尽量保证CPU环境的框架版本和GPU训练时完全一致
• 如果必须换版本，先测试一个简单的FCN模型（比如小尺寸输入）在CPU上是否正常，逐步定位问题

5. 验证模型导出的完整性 #

手动导出pb模型时，可能会遗漏某些关键变量（比如BatchNorm的滑动均值），导致CPU推理时用初始值计算出NaN。

解决办法：

• 放弃手动导出pb，改用框架原生的保存方法：比如TF的tf.saved_model.save，它会自动打包所有依赖的变量和计算图
• 加载模型后，打印模型的变量列表，确认moving_mean、moving_variance等关键变量的数值和GPU训练后的一致

先从浮点精度和算子差异这两点入手排查，这是最常见的原因，大概率能解决问题。

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