Comment obtenir les GPUs actuellement disponibles dans tensorflow ?

Question

J'ai l'intention d'utiliser TensorFlow distribué, et j'ai vu que TensorFlow peut utiliser les GPU pour la formation et les tests. Dans un environnement de cluster, chaque machine peut avoir 0 ou 1 ou plusieurs GPU, et je veux exécuter mon graphe TensorFlow dans les GPU sur autant de machines que possible.

J'ai découvert qu'en exécutant tf.Session() TensorFlow donne des informations sur le GPU dans les messages du journal comme ci-dessous :

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)

Ma question est la suivante : comment puis-je obtenir des informations sur le GPU actuellement disponible dans TensorFlow ? Je peux obtenir des informations sur les GPU chargés à partir du journal, mais je veux le faire d'une manière plus sophistiquée et programmatique. Je pourrais également restreindre les GPU intentionnellement en utilisant la variable d'environnement CUDA_VISIBLE_DEVICES, donc je ne veux pas connaître un moyen d'obtenir des informations sur les GPU à partir du noyau de l'OS.

En bref, je veux une fonction comme tf.get_available_gpus() qui renverra ['/gpu:0', '/gpu:1'] s'il y a deux GPU disponibles dans la machine. Comment puis-je mettre cela en œuvre ?

Answer 1

Il existe une méthode non documentée appelée device_lib.list_local_devices() qui vous permet de lister les dispositifs disponibles dans le processus local. ( N.B. En tant que méthode non documentée, elle est sujette à des modifications incompatibles avec le passé). La fonction renvoie une liste de DeviceAttributes tampon de protocole objets. Vous pouvez extraire une liste de noms de périphériques de type chaîne pour les périphériques GPU comme suit :

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

Notez que (au moins jusqu'à TensorFlow 1.4), appeler device_lib.list_local_devices() va exécuter un code d'initialisation qui, par défaut, va allouer toute la mémoire du GPU sur tous les périphériques ( Numéro de GitHub ). Pour éviter ce problème, créez d'abord une session avec une taille explicitement réduite de per_process_gpu_fraction o allow_growth=True pour éviter que toute la mémoire ne soit allouée. Voir cette question pour plus de détails.

Answer 2

Vous pouvez vérifier la liste de tous les appareils en utilisant le code suivant :

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

Answer 3

Il existe également une méthode dans le test utilitaire . Donc tout ce qu'il y a à faire est :

tf.test.is_gpu_available()

et/ou

tf.test.gpu_device_name()

Consultez la documentation de Tensorflow pour les arguments.

Answer 4

El réponse acceptée vous donne le nombre de GPU mais alloue également toute la mémoire sur ces GPU. Vous pouvez éviter cela en créant une session avec une mémoire inférieure fixe avant d'appeler device_lib.list_local_devices(), ce qui peut être indésirable pour certaines applications.

J'ai fini par utiliser nvidia-smi pour obtenir le nombre de GPU sans leur allouer de mémoire.

import subprocess

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

Answer 5

En dehors de l'excellente explication de Mrry, où il a suggéré d'utiliser device_lib.list_local_devices() Je peux vous montrer comment vérifier les informations relatives à la GPU à partir de la ligne de commande.

Comme actuellement seuls les gpus de Nvidia fonctionnent pour les frameworks NN, la réponse ne concerne qu'eux. Nvidia a une page où ils expliquent comment utiliser l'interface du système de fichiers /proc pour obtenir des informations d'exécution sur le pilote, les cartes graphiques NVIDIA installées et le statut AGP.

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

Fournir des informations sur chacune des cartes graphiques NVIDIA installées (nom du modèle, IRQ, version du BIOS, type de bus). type de bus). Notez que la version du BIOS n'est disponible que lorsque X est en cours d'exécution.

Vous pouvez donc exécuter cette commande depuis la ligne de commande cat /proc/driver/nvidia/gpus/0/information et voir les informations sur votre premier GPU. Il est facile de exécutez ceci à partir de python et vous pouvez aussi vérifier le deuxième, troisième, quatrième GPU jusqu'à ce qu'il échoue.

La réponse de Mrry est certainement plus robuste et je ne suis pas sûr que ma réponse fonctionnera sur une machine non-linux, mais la page de Nvidia fournit d'autres informations intéressantes, que peu de gens connaissent.