python - 在找到 GPU 后，如何强制 tensorflow 忘记它？

Question

我有一个基于 tensorflow 的代码，我在各种计算机上运行它，一些带有 CPU，一些带有 CPU 和 GPU。

如果机器上有 GPU，我想为用户提供使用 CPU 的选项。

来自 this answer 的代码工作正常:

import os
import tensorflow as tf

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

if tf.test.gpu_device_name():
    print('GPU found')
else:
    print("No GPU found")

# No GPU found

但是，我想检查 GPU 是否可用 第一 ，然后禁用它。

我试过:

import tensorflow as tf

if tf.test.gpu_device_name():
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

if tf.test.gpu_device_name():
    print('GPU found')

# GPU found

但它不起作用。一旦我使用 tf.test.gpu_device_name() ，它总是记得系统有一个 GPU。

我也试过 del tf , importlib.reload(tf) ，无济于事。唯一有效的是退出解释器并运行上面的第一个脚本。

一旦找到 GPU，如何让代码“忘记”它？

Answer 1

我不明白为什么你需要让 TensorFlow 忘记。你有 GPU 并不意味着你必须使用 GPU。

您可以使用 tf.device 指定底层设备。

例如:

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

c = tf.matmul(a, b)
print(c)

所以即使你有 GPU，程序仍然会使用 CPU。