[python] Comment obtenir les GPU disponibles actuellement dans tensorflow?


1 Answers

Vous pouvez vérifier toute la liste des appareils en utilisant le code suivant:

from tensorflow.python.client import device_lib

device_lib.list_local_devices()
Question

J'ai prévu d'utiliser TensorFlow distribué, et j'ai vu que TensorFlow peut utiliser des 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 graphique TensorFlow dans des GPU sur autant de machines que possible.

J'ai trouvé que lors de l'exécution de tf.Session() TensorFlow donne des informations sur 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 comment puis-je obtenir des informations sur le GPU disponible actuellement de TensorFlow? Je peux obtenir des informations GPU chargées à partir du journal, mais je veux le faire de manière plus sophistiquée et programmatique. Je pourrais également restreindre intentionnellement les GPU en utilisant la variable d'environnement CUDA_VISIBLE_DEVICES, donc je ne veux pas savoir comment obtenir des informations GPU à partir du noyau OS.

En bref, je veux une fonction comme tf.get_available_gpus() qui retournera ['/gpu:0', '/gpu:1'] s'il y a deux GPU disponibles dans la machine. Comment puis-je l'implémenter?




Il y a aussi une méthode dans l' utilitaire de test . Donc, tout ce qui doit être fait est:

tf.test.is_gpu_available()

et / ou

tf.test.gpu_device_name()

Recherchez les documents Tensorflow pour les arguments.




Related