Question

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

Answer 1

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

如果你有两个GPU并且Tensorflow同时看到它们，它将在两个设备上分配相同的内存，它将使用一些调度算法在需要时将元素从一个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)

Tensorflow默认使用哪种GPU

2 个答案: