Conozca lo último en aprendizaje automático, IA generativa y más en el Simposio WiML 2023.

Se usó la API de Cloud Translation para traducir esta página.

Usa una GPU

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar libreta

El código TensorFlow y los modelos tf.keras se ejecutarán de forma transparente en una sola GPU sin necesidad de cambios en el código.

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

Esta guía es para usuarios que probaron estos enfoques y descubrieron que necesitan un control detallado de cómo TensorFlow usa la GPU. Para obtener información sobre cómo depurar problemas de rendimiento para escenarios de GPU única y múltiple, consulte la guía Optimizar el rendimiento de GPU de TensorFlow .

Configuración

Asegúrate 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

Descripció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 su máquina.
"/GPU:0" : notación abreviada para la primera GPU de su máquina que es visible para TensorFlow.
"/job:localhost/replica:0/task:0/device:GPU:1" : nombre completo de la segunda GPU de su máquina que es visible para TensorFlow.

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

Si una operación de TensorFlow no tiene una implementación de GPU correspondiente, la operación vuelve al dispositivo de CPU. Por ejemplo, dado que tf.cast solo tiene un kernel de CPU, en un sistema con dispositivos CPU:0 y GPU:0 , el dispositivo CPU:0 se selecciona para ejecutar tf.cast , incluso si se solicita que se ejecute en el dispositivo GPU:0 .

Ubicación del dispositivo de registro

Para averiguar a qué dispositivos se asignan sus operaciones y tensores, coloque tf.debugging.set_log_device_placement(True) como la primera declaración de su programa. Habilitar el registro de colocación de dispositivos hace que se impriman 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 imprimirá una indicación de que la MatMul se ejecutó en GPU:0 .

Colocación manual de dispositivos

Si desea que una operación particular se ejecute en un dispositivo de su elección en lugar de lo que se selecciona automáticamente para usted, puede usar with tf.device para crear un contexto de dispositivo, y todas las operaciones dentro de ese contexto se ejecutarán 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 . Dado que no se especificó explícitamente un dispositivo para la operación de MatMul , el tiempo de ejecución de TensorFlow elegirá uno según la operación y los dispositivos disponibles ( GPU:0 en este ejemplo) y copiará automáticamente los tensores entre dispositivos si es necesario.

Limitar el crecimiento de la memoria de la GPU

De forma predeterminada, TensorFlow asigna casi toda la memoria de GPU de todas las GPU (sujetas a CUDA_VISIBLE_DEVICES ) visibles para el proceso. Esto se hace para utilizar de manera más eficiente los recursos de memoria GPU relativamente preciados en los dispositivos al reducir la fragmentación de la memoria. Para limitar TensorFlow a un conjunto específico de GPU, use el método tf.config.set_visible_devices .

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 la memoria según lo necesite el proceso. TensorFlow proporciona dos métodos para controlar esto.

La primera opción es activar el crecimiento de la memoria llamando a tf.config.experimental.set_memory_growth , que intenta asignar solo la cantidad de memoria de GPU necesaria para las asignaciones de tiempo de ejecución: comienza asignando muy poca memoria y, a medida que el programa se ejecuta y se necesita más memoria de GPU, la región de memoria de GPU se amplía para el proceso de TensorFlow. La memoria no se libera, ya que puede provocar la fragmentación de la memoria. Para activar el crecimiento de la memoria para una GPU específica, use el siguiente código antes de asignar cualquier tensor o ejecutar cualquier operación.

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 forma de habilitar esta opción es establecer la variable ambiental TF_FORCE_GPU_ALLOW_GROWTH en true . Esta configuración es específica de la plataforma.

El segundo método es configurar un dispositivo de GPU virtual con tf.config.set_logical_device_configuration y establecer un límite estricto 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 limitar 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.

Uso de 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 que TensorFlow elija automáticamente un dispositivo existente y compatible para ejecutar las operaciones en caso de que el especificado no exista, puede llamar a 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)

Uso de varias GPU

Desarrollar 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 probar fácilmente las 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 haya múltiples GPU lógicas disponibles para el tiempo de ejecución, puede utilizar las múltiples GPU con tf.distribute.Strategy o con ubicación manual.

Con `tf.distribute.Strategy`

La mejor práctica para usar varias GPU es usar tf.distribute.Strategy . Aquí hay un ejemplo simple:

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 ejecutará una copia de su modelo en cada GPU, dividiendo los datos de entrada entre ellos, también conocido como " paralelismo de datos ".

Para obtener más información sobre las estrategias de distribución, consulta la guía aquí .

Colocación manual

tf.distribute.Strategy funciona bajo el capó al replicar el cálculo en todos los dispositivos. Puede implementar manualmente la replicación 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)