TensorFlow脚本GPU初始化失败时如何抛出异常终止运行？

API 探路者

2026-5-21

强制TensorFlow在GPU初始化失败时抛出异常终止脚本

这问题我之前也碰到过，TensorFlow默认的自动 fallback到CPU的行为确实不符合咱们独占GPU场景的需求。想要让脚本在GPU初始化失败时直接终止而不是继续跑CPU，可以通过以下几个步骤实现：

1. 先检查系统是否有可用GPU

在脚本最开始先检测物理GPU是否存在，如果没有直接抛出异常：

import tensorflow as tf

# 获取所有可用的物理GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
if not gpus:
    raise RuntimeError("未检测到可用GPU，脚本终止运行")

2. 关闭TensorFlow的软设备放置功能

默认情况下，TensorFlow会在指定设备不可用时自动切换到其他可用设备（比如CPU），关闭这个功能后，当GPU不可用就会直接报错：

# 禁止自动 fallback到其他设备
tf.config.set_soft_device_placement(False)

3. 配置GPU内存（可选但推荐）

为了避免TensorFlow一次性占用全部GPU内存，同时确保内存分配失败时能抛出异常，可以设置内存按需分配：

try:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
    print(f"GPU内存配置出错: {e}")
    raise RuntimeError("GPU初始化失败，脚本终止运行")

4. 验证GPU可用性并触发异常

最后，强制在GPU上执行一个简单运算，确保GPU能正常初始化，如果失败就捕获异常并终止脚本：

try:
    # 强制指定使用第一个GPU
    with tf.device('/GPU:0'):
        # 执行一个简单的张量运算验证GPU是否可用
        test_tensor = tf.constant([1.0, 2.0]) + tf.constant([3.0, 4.0])
        print(f"GPU初始化成功，测试运算结果: {test_tensor.numpy()}")
except RuntimeError as e:
    raise RuntimeError("GPU初始化失败，脚本终止运行") from e