¿Cómo obtener las GPU disponibles actuales en tensorflow?


165

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?

Respuestas:


244

Hay un método no documentado llamado device_lib.list_local_devices()que le permite enumerar los dispositivos disponibles en el proceso local. ( Nota: como método no documentado, está sujeto a cambios incompatibles con versiones anteriores). La función devuelve una lista de objetos de DeviceAttributesbúfer de protocolo . Puede extraer una lista de nombres de dispositivos de cadena para los dispositivos GPU de la siguiente manera:

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

Tenga en cuenta que (al menos hasta TensorFlow 1.4), la llamada device_lib.list_local_devices()ejecutará un código de inicialización que, de forma predeterminada, asignará toda la memoria de la GPU en todos los dispositivos ( problema de GitHub ). Para evitar esto, primero cree una sesión con un tamaño explícitamente pequeño per_process_gpu_fractiono allow_growth=Truepara evitar que se asigne toda la memoria. Vea esta pregunta para más detalles.


12
PD: si este método alguna vez se mueve / cambia de nombre, miraría dentro de tensorflow / python / platform / test.py: is_gpu_available ya que se está utilizando bastante
Yaroslav Bulatov

1
¿Hay alguna forma de obtener los dispositivos de memoria libre y total? Veo que hay un campo memory_limit en DeviceAttributes y creo que es la memoria libre y no total
Aarbelle

2
Recuerdo que para versiones anteriores a 1 tensorflow imprimiría información sobre gpus cuando se importó en python. ¿Se han eliminado esos mensajes en las versiones más nuevas de tensorflow? (de ahí su sugerencia, la única forma de verificar las cosas de gpu)?
Charlie Parker el

@CharlieParker Creo que todavía imprimimos una línea de registro por dispositivo GPU al inicio en TF1.1.
Sr.

1
@aarbelle: el uso del método mencionado anteriormente para devolver todos los atributos incluye un campo Free memorypara mí, usando tensorflow1.1. En python:, from tensorflow.python.client import device_libentoncesdevice_lib.list_local_devices()
n1k31t4

123

Puede consultar la lista de todos los dispositivos con el siguiente código:

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

11
@Kulbear porque contiene estrictamente menos información que la respuesta existente.
Davidmh

3
Todavía prefiero esta respuesta debido a su simplicidad. Lo estoy usando directamente desde bash:python3 -c "from tensorflow.python.client import device_lib; print(device_lib.list_local_devices())"
aboettcher

1
Estoy de acuerdo, esta respuesta me ahorró tiempo. Acabo de copiar / pegar el código sin tener que leer la respuesta oficial más larga. Conozco los detalles, solo necesitaba la línea de código. Ya no fue elegida como respuesta y eso es suficiente. No hay necesidad de votar abajo.
Steven

1
obteniendo errorcannot import name 'format_exc' from 'traceback'
Siddharth Das

43

También hay un método en la utilidad de prueba . Entonces, todo lo que hay que hacer es:

tf.test.is_gpu_available()

y / o

tf.test.gpu_device_name()

Busque los documentos de Tensorflow para ver argumentos.


2
Esto devuelve solo GPU: 0
Trisoloriansunscreen

@Tal significa que tiene 1 GPU disponible (en la ranura PCI ID 0). Así tf.test.is_gpu_available()que volveremosTrue
respondió el

44
El OP solicitó un método que devuelve una lista de GPUS disponibles. Al menos en mi configuración de GPU múltiple, tf.test.gpu_device_name () devuelve solo el nombre del primero.
Trisoloriansunscreen

AttributeError: el módulo 'tensorflow' no tiene atributo 'test'
Siddharth Das

27

En TensorFlow 2.0, puede usar 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)

Si tiene dos GPU instaladas, genera esto:

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

Desde 2.1, puede soltar experimental:

gpus = tf.config.list_physical_devices('GPU')

Ver:


El comando funcionó muy bien. Tuve que cambiar 'GPU'a 'XLA_GPU'.
Vivek Subramanian

19

La respuesta aceptada le da la cantidad de GPU, pero también asigna toda la memoria en esas GPU. Puede evitar esto creando una sesión con memoria inferior fija antes de llamar a device_lib.list_local_devices () que puede ser no deseado para algunas aplicaciones.

Terminé usando nvidia-smi para obtener la cantidad de GPU sin asignarles memoria.

import subprocess

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

8

Además de la excelente explicación de Mrry, donde sugirió usarla device_lib.list_local_devices(), puedo mostrarle cómo puede verificar la información relacionada con la GPU desde la línea de comandos.

Debido a que actualmente solo los gpus de Nvidia funcionan para los marcos de NN, la respuesta solo los cubre. Nvidia tiene una página donde documentan cómo puede usar la interfaz del sistema de archivos / proc para obtener información de tiempo de ejecución sobre el controlador, las tarjetas gráficas NVIDIA instaladas y el estado de AGP.

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

Proporcione información sobre cada uno de los adaptadores gráficos NVIDIA instalados (nombre del modelo, IRQ, versión del BIOS, tipo de bus). Tenga en cuenta que la versión del BIOS solo está disponible mientras X se está ejecutando.

Entonces puede ejecutar esto desde la línea de comandos cat /proc/driver/nvidia/gpus/0/informationy ver información sobre su primera GPU. Es fácil ejecutar esto desde python y también puede verificar la segunda, tercera y cuarta GPU hasta que falle.

Definitivamente, la respuesta de Mrry es más sólida y no estoy seguro de si mi respuesta funcionará en una máquina que no sea Linux, pero esa página de Nvidia proporciona otra información interesante, que no mucha gente conoce.


3

Lo siguiente funciona en tensorflow 2:

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

Desde 2.1, puede soltar experimental:

    gpus = tf.config.list_physical_devices('GPU')

https://www.tensorflow.org/api_docs/python/tf/config/list_physical_devices


¿Funciona esto cuando también uso una scaleTier de BASIC_GPU? Cuando ejecuto este código, me da solo las CPUs
shivas

Respuesta duplicada de MiniQuark (pero con menos detalles ..)
FluxLemur

1

Tengo una GPU llamada NVIDIA GTX GeForce 1650 Ti en mi máquina contensorflow-gpu==2.2.0

Ejecute las siguientes dos líneas de código:

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

Salida:

Num GPUs Available:  1

0

Use esta forma y verifique todas las partes:

from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds


version = tf.__version__
executing_eagerly = tf.executing_eagerly()
hub_version = hub.__version__
available = tf.config.experimental.list_physical_devices("GPU")

print("Version: ", version)
print("Eager mode: ", executing_eagerly)
print("Hub Version: ", h_version)
print("GPU is", "available" if avai else "NOT AVAILABLE")

0

Asegúrese de tener instalada la última GPU TensorFlow 2.x en su máquina compatible con GPU, ejecute el siguiente código en python,

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf 

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

Obtendrá una salida como,

2020-02-07 10: 45: 37.587838: I tensorflow / stream_executor / cuda / cuda_gpu_executor.cc: 1006] la lectura exitosa del nodo NUMA de SysFS tuvo un valor negativo (-1), pero debe haber al menos un nodo NUMA, por lo que debe regresar NUMA node zero 2020-02-07 10: 45: 37.588896: I tensorflow / core / common_runtime / gpu / gpu_device.cc: 1746] Agregar dispositivos de gpu visibles: 0, 1, 2, 3, 4, 5, 6, 7 Num GPUs disponibles: 8


1
Esta respuesta ya ha sido publicada.
Björn Lindqvist
Al usar nuestro sitio, usted reconoce que ha leído y comprende nuestra Política de Cookies y Política de Privacidad.
Licensed under cc by-sa 3.0 with attribution required.