[英]Iterate cpu and gpu devices in Tensorflow
I am aware that Tensorflow can explicitly place computation on any devices by "/cpu0"
or "/gpu0"
. 我知道Tensorflow可以通过
"/cpu0"
或"/gpu0"
在任何设备上显式地放置计算。 However, this is hard-coded. 但是,这是硬编码的。 Is there any way to iterate all visible devices with built-in API?
有什么方法可以使用内置API迭代所有可见设备?
Here is what you would like to have: 这是您想要的:
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
Code is inspired from the best answer here: How to get current available GPUs in tensorflow? 此处的最佳答案启发了代码: 如何在tensorflow中获得当前可用的GPU?
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.