web-dev-qa-db-ja.com

テンソルフローで現在利用可能なGPUを取得する方法

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

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情報を取得する方法を知りたくありません。

手短に言えば、私はマシンに2つのGPUがある場合に['/gpu:0', '/gpu:1']を返すtf.get_available_gpus()のような関数が欲しいです。どうやってこれを実装できますか？

pythongputensorflow
99
Sangwon Kim

文書化されていない 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 issue ）。これを回避するには、最初に明示的に小さいper_process_gpu_fractionまたはallow_growth=Trueを使用してセッションを作成し、すべてのメモリが割り当てられないようにします。詳細については この質問 を参照してください。

180
mrry

次のコードを使ってすべてのデバイスリストを確認できます。

from tensorflow.python.client import device_lib

device_lib.list_local_devices()
85
hyun woo Cho

test util にもメソッドがあります。それで、やらなければならないことは次のとおりです。

tf.test.is_gpu_available()

および/または

tf.test.gpu_device_name()

引数についてはTensorflowのドキュメントを調べてください。

17
Soham Bhattacharyya

Mrryがdevice_lib.list_local_devices()を使うことを提案した素晴らしい説明とは別に、コマンドラインからGPU関連の情報をチェックする方法を紹介します。

現在NvidiaのgpusだけがNNフレームワークのために働くので、答えはそれらだけをカバーします。 Nvidiaには というページがあり、そこには/ procファイルシステムインタフェースを使ってドライバ、インストールされたNVIDIAグラフィックスカード、およびAGPステータスに関する実行時情報を取得する方法が記載されています。

/proc/driver/nvidia/gpus/0..N/information

インストールされている各NVIDIAグラフィックスアダプタに関する情報（モデル名、IRQ、BIOSバージョン、バスタイプ）を入力します。 BIOSのバージョンはXが動作している間だけ利用可能であることに注意してください。

そのため、これをコマンドラインcat /proc/driver/nvidia/gpus/0/informationから実行して、最初のGPUに関する情報を見ることができます。 pythonからこれを実行するのは簡単です また、失敗するまで2、3、4番目のGPUをチェックできます。

間違いなくMrryの答えはもっと堅牢で、私の答えがLinux以外のマシンでもうまくいくかどうかはわかりませんが、Nvidiaのページには他にも興味深い情報があり、あまり知られていません。

8
Salvador Dali

受け入れられた答え はあなたにGPUの数を与えますが、それはまたそれらのGPU上のすべてのメモリを割り当てます。これを回避するには、device_lib.list_local_devices（）を呼び出す前に低い固定メモリでセッションを作成します。これは一部のアプリケーションでは不要な場合があります。

メモリを割り当てずにGPUの数を取得するために、nvidia-smiを使用しました。

import subprocess

n = str(subprocess.check_output(["nvidia-smi", "-L"])).count('UUID')
8
mamad amin

TensorFlow 2.0では、tf.config.experimental.list_physical_devices('GPU')を使用できます。

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

2つのGPUがインストールされている場合、次のように出力されます。

Name: /physical_device:GPU:0   Type: GPU
Name: /physical_device:GPU:1   Type: GPU

詳細については このドキュメントページ をご覧ください。

3
MiniQuark