为啥 TensorFlow 总是使用 GPU 0？

Posted 2023-02-23

【中文标题】为啥 TensorFlow 总是使用 GPU 0？

【英文标题】：Why does TensorFlow always use GPU 0?为什么 TensorFlow 总是使用 GPU 0？

【发布时间】：2019-02-02 16:20:57

【问题描述】：

在多 GPU 设置上运行 TensorFlow 推理时遇到问题。

环境：Python 3.6.4； TensorFlow 1.8.0； Centos 7.3； 2 英伟达 Tesla P4

这是系统空闲时的 nvidia-smi 输出：

Tue Aug 28 10:47:42 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P4            Off  | 00000000:00:0C.0 Off |                    0 |
| N/A   38C    P0    22W /  75W |      0MiB /  7606MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla P4            Off  | 00000000:00:0D.0 Off |                    0 |
| N/A   39C    P0    23W /  75W |      0MiB /  7606MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

与我的问题相关的关键陈述：

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

def get_sess_and_tensor(ckpt_path):
    assert os.path.exists(ckpt_path), "file:  not exist.".format(ckpt_path)
    graph = tf.Graph()
    with graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(ckpt_path, "rb") as fid1:
            od_graph_def.ParseFromString(fid1.read())
            tf.import_graph_def(od_graph_def, name="")
        sess = tf.Session(graph=graph)
    with tf.device('/gpu:1'):
        tensor = graph.get_tensor_by_name("image_tensor:0")
        boxes = graph.get_tensor_by_name("detection_boxes:0")
        scores = graph.get_tensor_by_name("detection_scores:0")
        classes = graph.get_tensor_by_name('detection_classes:0')

    return sess, tensor, boxes, scores, classes

所以，问题是，当我将可见设备设置为“0,1”时，即使我将 tf.device 设置为 GPU 1，在运行推理时，我从 nvidia-smi 看到只使用了 GPU 0（GPU 0 的 GPU-Util 很高——几乎 100%——而 GPU 1 的为 0)。为什么不使用 GPU 1？

我想并行使用两个 GPU，但即使使用以下代码，它仍然只使用 GPU 0：

with tf.device('/gpu:0'):
    tensor = graph.get_tensor_by_name("image_tensor:0")
    boxes = graph.get_tensor_by_name("detection_boxes:0")
with tf.device('/gpu:1'):
    scores = graph.get_tensor_by_name("detection_scores:0")
    classes = graph.get_tensor_by_name('detection_classes:0')

非常感谢任何建议。

谢谢。

韦斯利

【问题讨论】：

至少你的设备名称不符合with the spec。应该是：/device:<type>:<num>。如果你使用/device:gpu:0,1 会发生什么？另见the note about the overridden device scope。

@agtoever 实际上，我在很多文章中都看到了 /gpu:0 格式，但不幸的是，我也尝试了你的建议，但遇到了同样的问题。

我认为您使用tf.device 为时已晚。您需要包装定义操作的代码。我不知道导入图形的东西会在哪里发生这种情况，但你可能想尝试移动 with tf.device 包装器，以便它包装 GraphDef 东西。

@xdurch0 你的意思是在训练过程中控制设备？实际上，对于训练，我没有指定可见设备和tf.device，但我认为它与这个问题无关，如果是这样，是否也意味着如果我在训练时指定GPU0，它只能使用GPU0进行推理?

我的意思是：你应该尝试类似with tf.device("/gpu:0"): tf.import_graph_def(...) 之类的东西。所以将device 的东西移到你要构建图表的位置。现在，您只能在从已经存在的图表中获取张量的部分周围使用它。这为时已晚；到那时，操作已经放在设备上（默认为 GPU0）。

【参考方案1】：

设备名称可能因您的设置而异。

执行：

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

并尝试将设备 name 用于您的第二个 GPU，与此处列出的完全相同。

【参考方案2】：

您可以使用 GPUtil 包来选择未使用的 gpus 并过滤 CUDA_VISIBLE_DEVICES 环境变量。

这将允许您在所有 GPU 上运行并行实验。

# Import os to set the environment variable CUDA_VISIBLE_DEVICES
import os
import tensorflow as tf
import GPUtil

# Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

# Get the first available GPU
DEVICE_ID_LIST = GPUtil.getFirstAvailable()
DEVICE_ID = DEVICE_ID_LIST[0] # grab first element from list

# Set CUDA_VISIBLE_DEVICES to mask out all other GPUs than the first available device id
os.environ["CUDA_VISIBLE_DEVICES"] = str(DEVICE_ID)

# Since all other GPUs are masked out, the first available GPU will now be identified as GPU:0
device = '/gpu:0'
print('Device ID (unmasked): ' + str(DEVICE_ID))
print('Device ID (masked): ' + str(0))

# Run a minimum working example on the selected GPU
# Start a session
with tf.Session() as sess:
    # Select the device
    with tf.device(device):
        # Declare two numbers and add them together in TensorFlow
        a = tf.constant(12)
        b = tf.constant(30)
        result = sess.run(a+b)
        print('a+b=' + str(result))

参考：https://github.com/anderskm/gputil