TensorFlow分配GPU显存但训练时利用率为0%，如何确认计算设备？

结合你遇到的场景（GPU显存已分配但利用率为0、训练速度和CPU版无差异），可以通过以下几种方式精准验证计算是否真的在GPU上执行：

1. 强制指定设备并对比耗时 #

直接将模型运算分别绑定到GPU和CPU，运行相同的小批量数据，对比两者的耗时差异——如果GPU有效运行，耗时会远低于CPU。

示例代码：

import tensorflow as tf
import time

# 准备测试用的小批量数据（匹配你的模型输入维度）
x_test = tf.random.normal((32, 224, 224, 3))  # 假设输入是224x224的RGB图像，批量32
y_test = tf.random.uniform((32, 10), maxval=10, dtype=tf.int32)

# 加载你的CNN模型
model = ...  # 替换为你的模型定义或加载代码

# 强制使用GPU运行
with tf.device('/GPU:0'):
    start = time.time()
    # 执行一次前向+反向传播（模拟训练步骤）
    with tf.GradientTape() as tape:
        preds = model(x_test, training=True)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_test, preds)
    grads = tape.gradient(loss, model.trainable_variables)
    print(f"GPU 单次训练步耗时: {time.time() - start:.4f}秒")

# 强制使用CPU运行
with tf.device('/CPU:0'):
    start = time.time()
    with tf.GradientTape() as tape:
        preds = model(x_test, training=True)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_test, preds)
    grads = tape.gradient(loss, model.trainable_variables)
    print(f"CPU 单次训练步耗时: {time.time() - start:.4f}秒")

如果GPU真的参与计算，耗时会比CPU少至少3-10倍；若两者耗时接近，说明GPU并未实际执行计算。

2. 解析设备分配日志细节 #

你已经启用了tf.debugging.set_log_device_placement(True)，需要重点关注核心计算操作的设备分配：

• 查找日志中类似Conv2D: (/job:localhost/replica:0/task:0/device:GPU:0)的条目，确认卷积、矩阵乘法、损失计算等关键步骤都分配到GPU
• 如果出现/device:CPU:0的核心计算操作，说明该步骤回退到CPU执行，这会导致整体速度无提升

3. 用NVIDIA Nsight Systems查看GPU活动 #

通过专业工具追踪GPU的实际运算行为，确认是否有CUDA kernel被调用：

1. 安装工具：

   sudo apt install nvidia-nsight-systems
   

2. 运行训练脚本并生成性能报告：

   nsys profile -t cuda,nvtx,osrt -o tf_train_profile python your_training_script.py
   

3. 打开生成的.qdrep文件（可用nsys-ui可视化），查看是否有cudnnConvolutionForward、cublasGemmEx等GPU计算kernel的执行记录。如果没有任何GPU kernel调用，说明计算完全在CPU上进行。

4. 验证张量运算的实际设备 #

检查训练过程中，模型参数、输入数据的实际存储设备，以及运算结果的设备：

import tensorflow as tf

# 检查GPU硬件是否被TensorFlow识别
print("已识别的GPU设备:", tf.config.list_physical_devices('GPU'))

# 检查模型参数的设备
print("模型第一层参数设备:", model.layers[0].kernel.device)

# 检查训练数据的设备
for x_batch, y_batch in train_dataset.take(1):
    print("输入数据设备:", x_batch.device)
    print("标签数据设备:", y_batch.device)

如果模型参数或输入数据仍在CPU上，即使TensorFlow分配了显存，实际计算也会回退到CPU。此时需要确保数据管道将数据加载到GPU，或手动转移：x_batch = x_batch.gpu()

补充排查点 #

结合你的情况（显存已分配但利用率低），还可以检查：

• 模型/批量大小：如果模型极小或批量太小，GPU的并行计算优势无法发挥，会出现利用率低、速度和CPU接近的情况，可尝试增大批量或模型规模
• XLA加速：开启XLA可能提升GPU利用率，在代码开头添加：

  tf.config.optimizer.set_jit(True)
  

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