Tengo un plan para usar TensorFlow distribuido, y vi que TensorFlow puede usar GPU para capacitación y pruebas. En un entorno de clúster, cada máquina podría tener 0 o 1 o más GPU, y quiero ejecutar mi gráfico TensorFlow en GPU en tantas máquinas como sea posible.
Descubrí que cuando se ejecuta tf.Session()
TensorFlow se brinda información sobre la GPU en los mensajes de registro como se muestra a continuación:
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)
Mi pregunta es ¿cómo obtengo información sobre la GPU actual disponible de TensorFlow? Puedo obtener información de GPU cargada del registro, pero quiero hacerlo de una manera más sofisticada y programática. También podría restringir las GPU intencionalmente usando la variable de entorno CUDA_VISIBLE_DEVICES, por lo que no quiero conocer una forma de obtener información de GPU del núcleo del sistema operativo.
En resumen, quiero que una función como tf.get_available_gpus()
esa regrese ['/gpu:0', '/gpu:1']
si hay dos GPU disponibles en la máquina. ¿Cómo puedo implementar esto?