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

Usa una GPU

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

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

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

Esta guía está dirigida a los usuarios que han probado estos enfoques y han descubierto que necesitan un control detallado de cómo TensorFlow usa la GPU.

Preparar

Asegúrate de tener instalada la última versión de TensorFlow gpu.

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

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

Si una operación de TensorFlow tiene implementaciones de CPU y GPU, de manera predeterminada, los dispositivos de GPU tendrán prioridad cuando la operación se asigne a un dispositivo. Por ejemplo, tf.matmul tiene núcleos de CPU y GPU. En un sistema con dispositivos CPU:0 y GPU:0 , el dispositivo GPU:0 se seleccionará para ejecutar tf.matmul menos que solicite explícitamente ejecutarlo en otro dispositivo.

Ubicación del dispositivo de registro

Para saber a qué dispositivos están asignados sus operaciones y tensores, coloque 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 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 la GPU:0 .

Colocación manual del dispositivo

Si desea que una operación en 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 un dispositivo no se especificó explícitamente para la operación MatMul , el tiempo de ejecución de TensorFlow elegirá uno en función de 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 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 GPU, usamos el método tf.config.experimental.set_visible_devices .

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.experimental.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)
2 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 activar el crecimiento de la memoria llamando a tf.config.experimental.set_memory_growth , que intenta asignar solo la cantidad de memoria 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 GPU, ampliamos la región de memoria GPU asignada al proceso de TensorFlow. Tenga en cuenta que no liberamos memoria, 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.experimental.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.experimental.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)
2 Physical GPUs, 2 Logical GPUs

Otra forma de habilitar esta opción es establecer la variable de entorno 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.experimental.set_virtual_device_configuration y establecer un límite tf.config.experimental.set_virtual_device_configuration en la memoria total para asignar en la GPU.

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.experimental.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)
2 Physical GPUs, 2 Logical GPUs

Esto es útil si desea vincular realmente 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)
/job:localhost/replica:0/task:0/device:GPU:2 unknown device.

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 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 desarrollamos en un sistema con una sola GPU, podemos simular múltiples GPU con dispositivos virtuales. Esto permite realizar pruebas sencillas de configuraciones de varias GPU sin necesidad de recursos adicionales.

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
         tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.experimental.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)
2 Physical GPU, 3 Logical GPUs

Una vez que tenemos múltiples GPU lógicas disponibles para el tiempo de ejecución, podemos utilizar las múltiples GPU con tf.distribute.Strategy o con colocación manual.

Con tf.distribute.Strategy

La mejor práctica para usar múltiples GPU es usar tf.distribute.Strategy . He aquí un ejemplo sencillo:

tf.debugging.set_log_device_placement(True)

strategy = tf.distribute.MirroredStrategy()
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 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 Add 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 VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Assert 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 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:1
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op VarIsInitializedOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op LogicalNot in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op Assert in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op Reshape 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 VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:1

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, consulte la guía aquí .

Colocación manual

tf.distribute.Strategy funciona bajo el capó replicando el cálculo en todos los 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.experimental.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 MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op AddN in device /job:localhost/replica:0/task:0/device:CPU:0
tf.Tensor(
[[ 44.  56.]
 [ 98. 128.]], shape=(2, 2), dtype=float32)