Use uma GPU

Veja no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

O código do TensorFlow e os modelos tf.keras serão executados de forma transparente em uma única GPU sem a necessidade de alterações de código.

A maneira mais simples de executar em várias GPUs, em uma ou várias máquinas, é usar Estratégias de Distribuição .

Este guia é para usuários que tentaram essas abordagens e descobriram que precisam de um controle refinado de como o TensorFlow usa a GPU. Para saber como depurar problemas de desempenho para cenários de GPU única e múltipla, consulte o guia Otimizar desempenho de GPU do TensorFlow .

Configurar

Verifique se você tem a versão mais recente do TensorFlow gpu instalada.

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

Visão geral

O TensorFlow suporta a execução de cálculos em vários tipos de dispositivos, incluindo CPU e GPU. Eles são representados com identificadores de string, por exemplo:

  • "/device:CPU:0" : A CPU da sua máquina.
  • "/GPU:0" : notação abreviada para a primeira GPU da sua máquina que é visível para o TensorFlow.
  • "/job:localhost/replica:0/task:0/device:GPU:1" : nome totalmente qualificado da segunda GPU da sua máquina que é visível para o TensorFlow.

Se uma operação do TensorFlow tiver implementações de CPU e GPU, por padrão, o dispositivo GPU será priorizado quando a operação for atribuída. Por exemplo, tf.matmul tem kernels de CPU e GPU e em um sistema com dispositivos CPU:0 e GPU:0 , o dispositivo GPU:0 é selecionado para executar tf.matmul , a menos que você solicite explicitamente para executá-lo em outro dispositivo.

Se uma operação do TensorFlow não tiver uma implementação de GPU correspondente, a operação retornará ao dispositivo da CPU. Por exemplo, como tf.cast possui apenas um kernel de CPU, em um sistema com dispositivos CPU:0 e GPU:0 , o dispositivo CPU:0 é selecionado para executar tf.cast , mesmo se solicitado a executar no dispositivo GPU:0 .

Posicionamento do dispositivo de registro

Para descobrir a quais dispositivos suas operações e tensores estão atribuídos, coloque tf.debugging.set_log_device_placement(True) como a primeira instrução do seu programa. A ativação do registro de posicionamento do dispositivo faz com que quaisquer alocações ou operações do Tensor sejam impressas.

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)

O código acima imprimirá uma indicação de que a MatMul foi executada em GPU:0 .

Posicionamento manual do dispositivo

Se você quiser que uma operação específica seja executada em um dispositivo de sua escolha em vez do que é selecionado automaticamente para você, você pode usar with tf.device para criar um contexto de dispositivo e todas as operações dentro desse contexto serão executadas no mesmo 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)

Você verá que agora a e b são atribuídos a CPU:0 . Como um dispositivo não foi especificado explicitamente para a operação MatMul , o runtime do TensorFlow escolherá um com base na operação e nos dispositivos disponíveis ( GPU:0 neste exemplo) e copiará automaticamente os tensores entre os dispositivos, se necessário.

Limitando o crescimento da memória da GPU

Por padrão, o TensorFlow mapeia quase toda a memória da GPU de todas as GPUs (sujeito a CUDA_VISIBLE_DEVICES ) visíveis para o processo. Isso é feito para usar com mais eficiência os recursos de memória da GPU relativamente preciosos nos dispositivos, reduzindo a fragmentação da memória. Para limitar o TensorFlow a um conjunto específico de GPUs, use o 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

Em alguns casos, é desejável que o processo aloque apenas um subconjunto da memória disponível ou apenas aumente o uso de memória conforme necessário pelo processo. O TensorFlow fornece dois métodos para controlar isso.

A primeira opção é ativar o crescimento de memória chamando tf.config.experimental.set_memory_growth , que tenta alocar apenas a memória GPU necessária para as alocações de tempo de execução: ele começa alocando muito pouca memória e, à medida que o programa é executado e mais memória da GPU é necessária, a região de memória da GPU é estendida para o processo do TensorFlow. A memória não é liberada, pois pode levar à fragmentação da memória. Para ativar o crescimento de memória para uma GPU específica, use o código a seguir antes de alocar qualquer tensor ou executar qualquer operação.

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

Outra maneira de habilitar essa opção é definir a variável de ambiente TF_FORCE_GPU_ALLOW_GROWTH como true . Essa configuração é específica da plataforma.

O segundo método é configurar um dispositivo de GPU virtual com tf.config.set_logical_device_configuration e definir um limite rígido na memória total a ser alocada na 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

Isso é útil se você quiser realmente limitar a quantidade de memória da GPU disponível para o processo do TensorFlow. Essa é uma prática comum para desenvolvimento local quando a GPU é compartilhada com outros aplicativos, como uma GUI de estação de trabalho.

Usando uma única GPU em um sistema multi-GPU

Se você tiver mais de uma GPU em seu sistema, a GPU com o ID mais baixo será selecionada por padrão. Se você quiser executar em uma GPU diferente, precisará especificar a preferência explicitamente:

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

Se o dispositivo que você especificou não existir, você receberá um RuntimeError : .../device:GPU:2 unknown device .

Se você quiser que o TensorFlow escolha automaticamente um dispositivo existente e compatível para executar as operações caso o especificado não exista, você pode chamar 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 várias GPUs

O desenvolvimento para várias GPUs permitirá que um modelo seja dimensionado com os recursos adicionais. Se estiver desenvolvendo em um sistema com uma única GPU, você pode simular várias GPUs com dispositivos virtuais. Isso permite o teste fácil de configurações de várias GPUs sem exigir recursos adicionais.

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

Quando houver várias GPUs lógicas disponíveis para o tempo de execução, você poderá utilizar as várias GPUs com tf.distribute.Strategy ou com posicionamento manual.

Com tf.distribute.Strategy

A melhor prática para usar várias GPUs é usar tf.distribute.Strategy . Aqui está um exemplo simples:

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 irá executar uma cópia do seu modelo em cada GPU, dividindo os dados de entrada entre elas, também conhecido como " paralelismo de dados ".

Para obter mais informações sobre estratégias de distribuição, consulte o guia aqui .

Colocação manual

tf.distribute.Strategy funciona nos bastidores, replicando a computação entre dispositivos. Você pode implementar manualmente a replicação construindo seu modelo em cada GPU. Por exemplo:

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)