在Tensorflow 1.14中，我试图使用tf.data.experimental.prefetch_to_device(device=...)将数据预取到GPU。但是我并不总是在GPU上训练，我经常在CPU上训练（尤其是在开发过程中）。
有没有办法让当前的默认设备使用？ Tensorflow要么选择CPU（当我设置CUDA_VISIBLE_DEVICES=-1时），否则它将选择GPU，默认情况下通常有效。
到目前为止，我只能找到一种方法来列出带有sess.list_devices()的可见设备，但是必须有一种查询当前默认设备的方法，因此我不必在prefetch_to_device中进行手动更改每次吧？
答案 0 :(得分：3)
目前没有API可以完成您所说的事情。最接近的
device = 'gpu:0' if tf.test.is_gpu_available() else 'cpu'
是您已经说过的。
我认为是因为分配是在较低级别完成的：https://github.com/tensorflow/tensorflow/blob/cf4dbb45ffb4d6ea0dc9c2ecfb514e874092cd16/tensorflow/core/common_runtime/colocation_graph.cc
也许您也可以尝试使用soft placement
希望有帮助。
答案 1 :(得分：3)
到目前为止，我发现的最佳解决方案是使用tf.test.is_gpu_available：
device = 'gpu:0' if tf.test.is_gpu_available() else 'cpu'