TensorFlow GPU版模型转CPU推理的nan异常线上部署方案

针对TensorFlow 2.2版本GPU推理正常、CPU环境同权重输出全NaN的问题，按操作优先级从高到低执行以下处理：

1. 关闭CPU端默认开启的高精度差异优化项 #

TensorFlow 2.2的CPU版本默认启用*oneDNN（MKL-DNN）*算子优化，该版本下部分算子（如BatchNorm、矩阵乘、Softmax）的数值实现与GPU端存在差异，尤其是对极小值、边界值的处理逻辑不同，极易触发数值发散导致NaN。
在导入TensorFlow的代码最前面添加全局配置，关闭相关优化：

import os
# 禁用oneDNN优化，消除跨端算子数值差异
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
# 关闭TensorFlow冗余日志，避免干扰排查
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import tensorflow as tf
# 强制全局使用32位浮点精度，禁用所有自动开启的混合精度配置
tf.keras.backend.set_floatx('float32')

2. 修正权重加载逻辑，避免层匹配错位 #

默认的model.load_weights()在跨平台加载时，容易因自定义层的初始化顺序差异，出现权重形状匹配但数值错位的问题，直接导致输出异常。替换原有加载逻辑为严格匹配模式：

model = Model()
# 按层名匹配权重，关闭不匹配自动跳过的逻辑，一旦加载错位直接抛出异常
model.load_weights('weights.hdf5', by_name=True, skip_mismatch=False)
# 加载完成后抽检首尾层的权重范围，确认无全0、全Inf的异常权重
for layer in model.layers[:2] + model.layers[-2:]:
    for w in layer.get_weights():
        print(f"{layer.name} 权重范围: [{w.min():.4f}, {w.max():.4f}]")

3. 补充边界值保护，消除除零类异常 #

GPU端推理时会在图优化阶段自动为除法类算子添加极小值epsilon保护，但TF2.2的CPU端没有对应的优化逻辑，遇到标准差为0、极小分母的场景会直接触发除零，输出NaN。

• 所有除法计算逻辑（如归一化、特征缩放）手动添加epsilon保护，不依赖框架默认实现：

# 原逻辑存在除零风险：y = (x - mean) / std
# 替换为带保护的逻辑
std_safe = std + 1e-7
y = (x - mean) / std_safe

• 输入模型前对数据做异常值清洗，避免带异常的输入触发算子发散：

import numpy as np
# 将输入中的NaN、Inf值替换为安全范围数值
input_data = np.nan_to_num(input_data, nan=0.0, posinf=1e6, neginf=-1e6)

4. 跨平台加载的兜底方案 #

如果以上步骤执行后仍有问题，直接在GPU环境下将模型导出为SavedModel格式，CPU端直接加载完整模型，规避hdf5权重跨平台加载的维度错位问题：

# 第一步：在GPU正常推理的环境中执行，导出完整模型
model = Model()
model.load_weights('weights.hdf5')
model.save('saved_model_output', save_format='tf')

# 第二步：在CPU服务环境中直接加载导出的完整模型开展推理
loaded_model = tf.keras.models.load_model('saved_model_output')
predict_result = loaded_model.predict(input_data)

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