¡El Día de la Comunidad de ML es el 9 de noviembre! Únase a nosotros para recibir actualizaciones de TensorFlow, JAX, y más Más información

Usa una GPU

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

Código TensorFlow y tf.keras modelos transparente funcionar con una sola GPU sin cambios de código necesarios.

La forma más sencilla de ejecutar en múltiples GPU, en una o varias máquinas, es el uso de estrategias de distribución .

Esta guía es para usuarios que han probado estos enfoques y han descubierto que necesitan un control detallado de cómo TensorFlow usa la GPU. Para aprender cómo depurar problemas de rendimiento para los escenarios individuales y multi-GPU, ver el Optimizar el rendimiento de la GPU TensorFlow guía.

Configuración

Asegúrese de tener instalada la última versión de GPU de TensorFlow.

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

Visión general

TensorFlow admite la ejecución de cálculos en una variedad de tipos de dispositivos, incluidos CPU y GPU. Se representan con identificadores de cadena, por ejemplo:

  • "/device:CPU:0" : La CPU de la máquina.
  • "/GPU:0" : la notación corta a mano por primera GPU de la máquina que es visible para TensorFlow.
  • "/job:localhost/replica:0/task:0/device:GPU:1" : Nombre completo de la segunda GPU de la máquina que es visible para TensorFlow.

Si una operación de TensorFlow tiene implementaciones de CPU y GPU, de forma predeterminada, el dispositivo GPU tiene prioridad cuando se asigna la operación. Por ejemplo, tf.matmul tiene tanto la CPU y la GPU y granos en un sistema con dispositivos CPU:0 y GPU:0 , la GPU:0 se selecciona el dispositivo para funcionar tf.matmul a menos que solicite explícitamente a ejecutarlo en otro dispositivo.

Si una operación de TensorFlow no tiene una implementación de GPU correspondiente, la operación recurre al dispositivo de la CPU. Por ejemplo, desde tf.cast sólo tiene un núcleo CPU, en un sistema con dispositivos de CPU:0 y GPU:0 , la CPU:0 dispositivo está seleccionado para ejecutarse tf.cast , incluso si solicitado para ejecutar en el GPU:0 dispositivo .

Ubicación del dispositivo de registro

Para saber qué dispositivos tensores sus operaciones y se asignan a, puesto tf.debugging.set_log_device_placement(True) como la primera declaración de su programa. Al habilitar el registro de ubicación de dispositivos, se imprimen las asignaciones u operaciones de Tensor.

tf.debugging.set_log_device_placement(True)

# Create some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

El código anterior se imprimirá una indicación del MatMul op fue ejecutado el GPU:0 .

Colocación manual del dispositivo

Si desea una operación en particular para ejecutarse en un dispositivo de su elección en lugar de lo que está seleccionada de forma automática para usted, usted puede utilizar with tf.device para crear un contexto de dispositivo, y todas las operaciones dentro de ese contexto se quedará en el mismo dispositivo designado .

tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Verá que ahora a y b están asignados a CPU:0 . Puesto que un dispositivo no se especifica explícitamente para el MatMul operación, el tiempo de ejecución TensorFlow elegirá uno basado en la operación y dispositivos disponibles ( GPU:0 en este ejemplo) y copiar automáticamente tensores entre dispositivos si es necesario.

Limitar el crecimiento de la memoria de la GPU

Por defecto, TensorFlow los mapas de casi toda la memoria de la GPU de todas las GPU (sujeto a CUDA_VISIBLE_DEVICES ) visibles para el proceso. Esto se hace para utilizar de manera más eficiente los recursos de memoria de la GPU relativamente valiosos en los dispositivos al reducir la fragmentación de la memoria. Para limitar TensorFlow a un conjunto específico de las GPU, utilice el tf.config.set_visible_devices método.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
  except RuntimeError as e:
    # Visible devices must be set before GPUs have been initialized
    print(e)
1 Physical GPUs, 1 Logical GPU

En algunos casos, es deseable que el proceso solo asigne un subconjunto de la memoria disponible, o que solo aumente el uso de memoria según lo necesite el proceso. TensorFlow proporciona dos métodos para controlar esto.

La primera opción es para encender el crecimiento de memoria llamando tf.config.experimental.set_memory_growth , que intenta asignar sólo la cantidad de memoria de la GPU según sea necesario para las asignaciones de tiempo de ejecución: comienza a cabo la asignación de memoria muy pequeña, y que el programa se ejecute y se necesita más memoria GPU, la región de memoria GPU se extiende para el proceso TensorFlow. La memoria no se libera ya que puede provocar su fragmentación. Para activar el crecimiento de la memoria para una GPU específica, use el siguiente código antes de asignar tensores o ejecutar operaciones.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)
Physical devices cannot be modified after being initialized

Otra manera de activar esta opción es establecer la variable ambiental TF_FORCE_GPU_ALLOW_GROWTH de true . Esta configuración es específica de la plataforma.

El segundo método es para configurar un dispositivo GPU virtual con tf.config.set_logical_device_configuration y establecer un límite duro en la memoria total para asignar en la GPU.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
Virtual devices cannot be modified after being initialized

Esto es útil si realmente desea vincular la cantidad de memoria de GPU disponible para el proceso de TensorFlow. Esta es una práctica común para el desarrollo local cuando la GPU se comparte con otras aplicaciones, como la GUI de una estación de trabajo.

Usando una sola GPU en un sistema multi-GPU

Si tiene más de una GPU en su sistema, la GPU con la ID más baja se seleccionará de forma predeterminada. Si desea ejecutar en una GPU diferente, deberá especificar la preferencia explícitamente:

tf.debugging.set_log_device_placement(True)

try:
  # Specify an invalid GPU device
  with tf.device('/device:GPU:2'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)
except RuntimeError as e:
  print(e)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0

Si el dispositivo que ha especificado no existe, obtendrá un RuntimeError : .../device:GPU:2 unknown device .

Si desea TensorFlow para elegir automáticamente un dispositivo existente y con el apoyo para ejecutar las operaciones en caso de que la especificada no existe, se puede llamar tf.config.set_soft_device_placement(True) .

tf.config.set_soft_device_placement(True)
tf.debugging.set_log_device_placement(True)

# Creates some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Usando múltiples GPU

El desarrollo para múltiples GPU permitirá que un modelo se escale con los recursos adicionales. Si desarrolla en un sistema con una sola GPU, puede simular varias GPU con dispositivos virtuales. Esto permite realizar pruebas sencillas de configuraciones de varias GPU sin necesidad de recursos adicionales.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024),
         tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
Virtual devices cannot be modified after being initialized

Una vez que hay múltiples GPU lógicos disponibles para el tiempo de ejecución, puede utilizar las múltiples GPU con tf.distribute.Strategy o con la colocación manual.

con tf.distribute.Strategy

La mejor práctica para el uso de múltiples GPU es utilizar tf.distribute.Strategy . A continuación, se muestra un ejemplo sencillo:

tf.debugging.set_log_device_placement(True)
gpus = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
  inputs = tf.keras.layers.Input(shape=(1,))
  predictions = tf.keras.layers.Dense(1)(inputs)
  model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
  model.compile(loss='mse',
                optimizer=tf.keras.optimizers.SGD(learning_rate=0.2))
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op RandomUniform in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Sub in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0

Este programa se ejecutará una copia de su modelo en cada GPU, la división de los datos de entrada entre ellas, también conocidos como " paralelismo de datos ".

Para obtener más información acerca de las estrategias de distribución, echa un vistazo a la guía aquí .

Colocación manual

tf.distribute.Strategy funciona bajo el capó mediante la replicación de cálculo a través de dispositivos. Puede implementar la replicación manualmente construyendo su modelo en cada GPU. Por ejemplo:

tf.debugging.set_log_device_placement(True)

gpus = tf.config.list_logical_devices('GPU')
if gpus:
  # Replicate your computation on multiple GPUs
  c = []
  for gpu in gpus:
    with tf.device(gpu.name):
      a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
      c.append(tf.matmul(a, b))

  with tf.device('/CPU:0'):
    matmul_sum = tf.add_n(c)

  print(matmul_sum)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)