2016-07-25 32 views
39

Ich habe einen Plan, um TensorFlow zu verwenden, und ich sah TensorFlow kann GPUs für Training und Tests verwenden. In einer Cluster-Umgebung kann jede Maschine 0 oder 1 oder mehr GPUs haben, und ich möchte meine TensorFlow-Grafik auf möglichst vielen Maschinen in GPUs umwandeln.Wie erhält man die aktuell verfügbaren GPUs im Tensorflow?

Ich fand, dass, wenn tf.Session() läuft TensorFlow in den Log-Meldungen wie folgenden Informationen über GPU gibt:

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) 

Meine Frage ist, wie kann ich über die aktuellen verfügbaren GPU Informationen erhalten von TensorFlow? Ich kann geladene GPU-Informationen aus dem Protokoll abrufen, aber ich möchte das auf eine anspruchsvollere, programmatische Weise tun. Ich könnte auch GPUs absichtlich mit der Umgebungsvariablen CUDA_VISIBLE_DEVICES einschränken, also möchte ich keine Möglichkeit kennen, GPU-Informationen vom Betriebssystemkern zu bekommen.

Kurz gesagt, ich möchte eine Funktion wie tf.get_available_gpus(), die ['/gpu:0', '/gpu:1'] zurückgibt, wenn zwei GPUs in der Maschine verfügbar sind. Wie kann ich das umsetzen?

Antwort

88

Es gibt eine undokumentierte Methode namens device_lib.list_local_devices(), mit der Sie die im lokalen Prozess verfügbaren Geräte auflisten können. (N.B. Als undokumentierte Methode unterliegt dies rückwärtskompatiblen Änderungen.) Die Funktion gibt eine Liste von DeviceAttributes protocol buffer Objekten zurück. Sie können für die GPU-Geräte eine Liste von String-Gerätenamen extrahieren wie folgt:

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'] 

Beachten Sie, dass (zumindest bis zu TensorFlow 1.4), device_lib.list_local_devices() Aufruf wird einige Initialisierungscode, die standardmäßig ausgeführt werden, werden alle zuteilen den GPU-Speicher auf allen Geräten (GitHub issue). Um dies zu vermeiden, erstellen Sie zunächst eine Sitzung mit einem explizit kleinen Wert per_process_gpu_fraction oder allow_growth=True, um zu verhindern, dass der gesamte Speicher zugeordnet wird. Weitere Informationen finden Sie unter this question.

+6

PS, wenn diese Methode jemals verschoben/umbenannt wird, würde ich in tensorflow/python/plattform/test.py suchen: is_gpu_available, da das ziemlich oft verwendet wird –

+1

Gibt es eine Möglichkeit, die Geräte frei und insgesamt Speicher zu bekommen? Ich sehe, dass es ein memory_limit Feld in den DeviceAttributes gibt und ich denke, es ist der freie Speicher und nicht – aarbelle

+0

Ich erinnere mich, dass für frühere Versionen als 1 Tensorflow einige Informationen über gpus drucken würde, wenn es in Python importiert wurde. Wurden diese Nachrichten in den neueren Tensorflow-Versionen entfernt? (daher Ihr Vorschlag der einzige Weg, um GPU Zeug zu überprüfen)? –

23

können Sie alle Geräteliste überprüfen folgenden Code:

from tensorflow.python.client import device_lib 

device_lib.list_local_devices() 
+0

Warum wurde dies downvoted ... – Kulbear

+2

@Kulbear, weil es streng weniger Informationen als die vorhandene Antwort enthält. – Davidmh

5

Neben der hervorragenden Erklärung von Mrry, wo er vorgeschlagen device_lib.list_local_devices() verwenden ich kann Ihnen zeigen, wie Sie für GPU bezogene Informationen aus dem Befehl überprüfen Linie.

Weil derzeit nur Nvidias GPUs für NN-Frameworks funktionieren, deckt die Antwort nur sie ab. Nvidia has a page Hier dokumentieren sie, wie Sie mit der Dateisystemschnittstelle/proc Laufzeitinformationen über den Treiber, alle installierten NVIDIA-Grafikkarten und den AGP-Status erhalten können.

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

Informationen über jeder der installierten NVIDIA-Grafikadapter bieten (Modell, IRQ, BIOS Version, Bustyp). Beachten Sie, dass die BIOS-Version nur verfügbar ist, während X ausgeführt wird.

So können Sie dies von der Befehlszeile cat /proc/driver/nvidia/gpus/0/information ausführen und Informationen über Ihre erste GPU sehen. Es ist einfach, run this from python und auch Sie können zweite, dritte, vierte GPU überprüfen, bis es fehlschlägt.

Auf jeden Fall ist Mrrys Antwort robuster und ich bin mir nicht sicher, ob meine Antwort auf Nicht-Linux-Rechnern funktioniert, aber dass Nvidias Seite andere interessante Informationen liefert, von denen nicht viele wissen.