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。

找到后，如何迫使张量流忘记GPU？

1 个答案: