在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

代码的灵感来自最佳答案：How to get current available GPUs in tensorflow?