在 Tensorflow 中迭代 cpu 和 gpu 设备

【问题标题】：Iterate cpu and gpu devices in Tensorflow在 Tensorflow 中迭代 cpu 和 gpu 设备
【发布时间】：2017-07-29 08:06:42
【问题描述】：

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

【问题讨论】：

How to get current available GPUs in tensorflow?的可能重复

标签： tensorflow tensorflow-gpu

【解决方案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?

【讨论】：