获取 GPU 设备名称 TensorFlow - Python

当使用 TensorFlow 运行计算时，可以通过以下代码片段获取 GPU 设备名称：

import tensorflow as tf

# 创建 TensorFlow 会话
sess = tf.Session()

# 使用 TensorFlow 输出 GPU 设备的名称
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU 设备未发现')
print('GPU 找到: {}'.format(device_name))

这段代码将创建一个 TensorFlow 会话并输出 GPU 设备名称。

如果找不到 GPU 设备，则会引发 SystemError。

请注意，TensorFlow 运行时需要正确安装 CUDA 和 cuDNN。

其他注意事项

• 如果有多个 GPU 设备，则可以使用以下代码选择要使用的设备：

  import tensorflow as tf
  
  # 创建 TensorFlow 会话并选择第二个 GPU 设备
  with tf.device('/device:GPU:1'):
      sess = tf.Session()
  

• 可以使用以下代码检查 TensorFlow 是否正在使用 GPU 运行：

  import tensorflow as tf
  
  # 创建 TensorFlow 会话并输出 TensorFlow 是否正在运行 GPU
  sess = tf.Session()
  
  # `device_lib.list_local_devices` 可以获取所有可用设备的列表，
  # 它将输出包含 GPU 和 CPU 的所有设备，因此需要适当过滤。
  device_list = [x.name for x in tf.device_lib.list_local_devices() if x.device_type == 'GPU']
  if len(device_list) > 0:
      print('TensorFlow 正在使用 GPU 运行')
      print('GPU 设备列表: {}'.format(device_list))
  else:
      print('TensorFlow 正在使用 CPU 运行')
  

• 强烈建议在运行 TensorFlow 时使用 GPU，因为它可以显著提高计算性能。

现在，您应该可以轻松地获取 GPU 设备名称，然后根据需要在 TensorFlow 代码中使用它了。