Tensorflow 獲取默認設備名稱

Question

在 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中手動更改它，對吧？

Answer 1

目前沒有 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

希望能幫助到你。

Answer 2

到目前為止我發現的最佳解決方案是使用tf.test.is_gpu_available ：

device = 'gpu:0' if tf.test.is_gpu_available() else 'cpu'