Tensorflow默认使用哪个GPU

Question

我的计算机上安装了两个GPU。当我运行我的代码时，我总是看到来自Tensorflow的调试信息，它将显示我的两个GPU卡的信息。然而，如果我没有指定TF，会使用哪一个？如果我有多个GPU卡，我是否需要为TF指定一个GPU以提高效率。谢谢~！

Answer 1

Tensorflow使用系统上所有可用的GPU。

如果你有两个GPU，Tensorflow同时看到两个GPU，它将在两个设备上分配相同的内存，并在需要时使用某种调度算法将元素从一个GPU内存交换到另一个GPU内存。

您可以使用with tf.device("/gpu:<number>")上下文管理器决定在哪个GPU中执行指定的操作。然而，内存是在每个GPU中分配的。

如果您希望只使用一些GPU，并避免在任何其他GPU上分配任何内容，则必须使用CUDA环境变量CUDA_VISIBLE_DEVICES。

在执行python脚本之前，您可以指定Tensorflow将看到的CPU。这样，您就可以确定只会使用某些设备。

例如，如果你想让Tensorflow只看到你的第二个GPU，你可以这样运行你的脚本：

CUDA_VISIBLE_DEVICES=1 python execute_script.py

Answer 2

如果你不想在每次运行代码之前添加export CUDA_VISIBLE_DEVICES=0或CUDA_VISIBLE_DEVICES=0 python xx.py，那么在Python代码中设置GPU选项可能会更好：

# Set GPU
os.environ['CUDA_VISIBLE_DEVICES'] = flags.gpu_choice

flags.gpu_choice可以是类似于'0,1,2'的字符串

然后，您可以简单地按照Use GPUs上的Tensorflow教程进行操作：

with tf.device('/device:GPU:2'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
  b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
  c = tf.matmul(a, b)