python - 如何检查 GPU 是否存在而不在 tensorflow 中生成警告

Question

• 我正在编写代码来训练模型，我不确定这段代码是否会在以后有或没有 GPU 的机器上运行，所以我正在使用代码。

• 我的问题是：代码正在生成此警告：

  • is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version. Instructions for updating: Use tf.config.list_physical_devices('GPU') instead.

所以我的问题是：

如何以简单明了的方式检查是否有 GPU，而不会产生警告。

import os

# Check if there is valid GPU to be used
if(tf.test.is_gpu_available(cuda_only=True)):
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'

标签： pythonpython-3.xtensorflowgpu

Answer 1

尝试使用list_physical_devices您的警告告诉您使用的 from tf.config.experimental 。这段代码对我有用。

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  print('GPU is available')