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

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'

Answer 1

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

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