TensorFlow 2.0加载模型无法分配至GPU的问题求助

针对你遇到的加载后模型无法分配到GPU的问题，我在类似场景下也碰到过，结合你的代码和日志信息，下面给出几个可行的解决方案：

1. 切换到SavedModel格式保存/加载模型 #

.h5格式虽然方便，但它主要聚焦于保存模型结构和权重，对于TensorFlow 2.x的动态图执行、设备布局等信息的保存不够完整。而TensorFlow原生的SavedModel格式会完整保存模型的计算图、配置和设备相关元数据，能更好地维持模型的运行环境。

修改你的代码如下：

# 保存模型时改用SavedModel格式
model.save("./saved_model")

# 加载模型
loaded_model = tf.keras.models.load_model("./saved_model")

# 运行预测（tf.device包裹可以保留，也可以省略，TF会自动使用可用GPU）
with tf.device("/GPU:0"):
    loaded_model.predict(...)

2. 加载后手动将模型转移到GPU #

如果必须使用.h5格式，加载模型后可以通过to()方法强制将整个模型的权重和计算图转移到GPU设备上：

loaded_model = tf.keras.models.load_model("./mod.h5")
# 将模型转移到GPU
loaded_model = loaded_model.to("/GPU:0")

# 直接运行预测即可，无需额外tf.device包裹
loaded_model.predict(...)

3. 确保GPU设备被TensorFlow正确识别 #

在加载模型前，可以先确认GPU是否被TF检测到，并强制设置GPU为可见设备：

# 检查可用GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    # 只使用第一个GPU设备
    tf.config.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(f"物理GPU: {len(gpus)}, 逻辑GPU: {len(logical_gpus)}")

# 之后再加载模型并运行
loaded_model = tf.keras.models.load_model("./mod.h5")
loaded_model = loaded_model.to("/GPU:0")
loaded_model.predict(...)

为什么会出现日志显示GPU分配但实际跑CPU的情况？ #

你的日志里placer确实显示将节点分配到了GPU，但实际执行时 fallback 到CPU，大概率是因为.h5模型加载后，权重默认留在了CPU内存中，TF在执行时发现权重不在GPU上，就自动切换到CPU运行（虽然日志里的初始分配是GPU）。通过上述方法将模型整体转移到GPU后，权重和计算图都会在GPU上，就能正常加速了。

另外针对Jetson TX2平台，要确保你的JetPack版本、CUDA/cuDNN版本和TensorFlow 2.x版本完全兼容，这也是GPU加速正常运行的前提。

内容的提问来源于stack exchange，提问作者Austin G-Dragon