tensorflow - 在 Tensorflow 中迭代 cpu 和 gpu 设备

Question

我知道 Tensorflow 可以通过"/cpu0"或明确地将计算放置在任何设备上"/gpu0"。但是，这是硬编码的。有没有办法使用内置 API 迭代所有可见设备？

Answer 1

这是您想要的：

import tensorflow as tf
from tensorflow.python.client import device_lib

def get_all_devices():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos]

all_devices = get_all_devices()
for device_name in all_devices:
    with tf.device(device_name):
        if "cpu" in device_name:
            # Do something
            pass
        if "gpu" in device_name:
            # Do something else
            pass

代码的灵感来自这里的最佳答案：如何在 tensorflow 中获取当前可用的 GPU？