[解決済み】tensorflowで現在利用可能なGPUを取得する方法とは？

質問

分散型TensorFlowを使用する計画があり、TensorFlowは学習とテストにGPUを使用できることを知りました。クラスタ環境では、各マシンは0または1以上のGPUを持つことができ、私はできるだけ多くのマシンのGPUにTensorFlowグラフを実行したいです。

を実行したところ tf.Session() TensorFlowは以下のようなログメッセージでGPUに関する情報を提供します。

I tensorflow/core/common_runtime/gpu/gpu_init.cc:126] DMA: 0 
I tensorflow/core/common_runtime/gpu/gpu_init.cc:136] 0:   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:838] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:01:00.0)

私の質問は、TensorFlowから現在利用可能なGPUの情報を得るにはどうしたらよいかということです。ログからロードされたGPUの情報を得ることはできますが、私はより洗練された、プログラム的な方法でそれをしたいです。 また、CUDA_VISIBLE_DEVICES環境変数を使って意図的にGPUを制限することもできるので、OSカーネルからGPU情報を取得する方法は知りたくありません。

要するに、次のような関数が欲しいのです。 tf.get_available_gpus() を返します。 ['/gpu:0', '/gpu:1'] マシンに利用可能なGPUが2つある場合。どのように実装すればよいのでしょうか？

どのように解決するのですか？

という文書化されていないメソッドがあります。 device_lib.list_local_devices() を使用すると、ローカルプロセスで使用可能なデバイスを一覧表示することができます。( N.B. 文書化されていないメソッドであるため、後方互換性のない変更が行われる可能性があります)。この関数は DeviceAttributes プロトコルバッファ オブジェクトを作成します。GPUデバイスの文字列デバイス名の一覧は、以下のように抽出できます。

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']

少なくともTensorFlow 1.4まで）。 device_lib.list_local_devices() はいくつかの初期化コードを実行し、デフォルトでは、すべてのデバイスにすべての GPU メモリを割り当てます ( GitHubの課題 ). これを避けるには、まず、明示的に小さな per_process_gpu_fraction または allow_growth=True を使用すると、すべてのメモリが確保されるのを防ぐことができます。参照 この質問 をご覧ください。