Precisão mista

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

Visão geral

A precisão mista é o uso de tipos de ponto flutuante de 16 bits e 32 bits em um modelo durante o treinamento para torná-lo mais rápido e usar menos memória. Ao manter certas partes do modelo nos tipos de 32 bits para estabilidade numérica, o modelo terá um tempo de etapa menor e treinará igualmente em termos de métricas de avaliação, como precisão. Este guia descreve como usar a API de precisão mista Keras para acelerar seus modelos. O uso dessa API pode melhorar o desempenho em mais de 3 vezes em GPUs modernas e 60% em TPUs.

Hoje, a maioria dos modelos usa o dtype float32, que consome 32 bits de memória. No entanto, existem dois dtypes de menor precisão, float16 e bfloat16, cada um que ocupa 16 bits de memória. Aceleradores modernos podem executar operações mais rapidamente nos dtypes de 16 bits, pois possuem hardware especializado para executar cálculos de 16 bits e os dtypes de 16 bits podem ser lidos da memória mais rapidamente.

As GPUs NVIDIA podem executar operações em float16 mais rápido que em float32, e as TPUs podem executar operações em bfloat16 mais rápido que float32. Portanto, esses dtypes de baixa precisão devem ser usados ​​sempre que possível nesses dispositivos. No entanto, variáveis ​​e alguns cálculos ainda devem estar em float32 por razões numéricas para que o modelo seja treinado com a mesma qualidade. A API de precisão mista Keras permite que você use uma combinação de float16 ou bfloat16 com float32, para obter os benefícios de desempenho de float16/bfloat16 e os benefícios de estabilidade numérica de float32.

Configurar

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision

Hardware compatível

Embora a precisão mista seja executada na maioria dos hardwares, ela apenas acelerará os modelos em GPUs NVIDIA e Cloud TPUs recentes. As GPUs NVIDIA suportam uma combinação de float16 e float32, enquanto as TPUs suportam uma combinação de bfloat16 e float32.

Entre as GPUs NVIDIA, aquelas com capacidade de computação 7.0 ou superior terão o maior benefício de desempenho da precisão mista porque possuem unidades de hardware especiais, chamadas Tensor Cores, para acelerar as multiplicações e convoluções da matriz float16. As GPUs mais antigas não oferecem nenhum benefício de desempenho matemático ao usar precisão mista, no entanto, a economia de memória e largura de banda pode permitir alguns aumentos de velocidade. Você pode pesquisar a capacidade de computação da sua GPU na página da Web da GPU CUDA da NVIDIA . Exemplos de GPUs que mais se beneficiarão da precisão mista incluem GPUs RTX, V100 e A100.

Você pode verificar seu tipo de GPU com o seguinte. O comando só existe se os drivers NVIDIA estiverem instalados, portanto, o seguinte gerará um erro.

nvidia-smi -L

GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-99e10c4d-de77-42ee-4524-6c41c4e5e47d)

Todos os Cloud TPUs são compatíveis com bfloat16.

Mesmo em CPUs e GPUs mais antigas, onde nenhuma aceleração é esperada, APIs de precisão mista ainda podem ser usadas para teste de unidade, depuração ou apenas para experimentar a API. No entanto, em CPUs, a precisão mista será significativamente mais lenta.

Configurando a política dtype

Para usar a precisão mista no Keras, você precisa criar um tf.keras.mixed_precision.Policy , normalmente chamado de política dtype . As políticas Dtype especificam as camadas dtypes em que serão executadas. Neste guia, você construirá uma política a partir da string 'mixed_float16' e a definirá como a política global. Isso fará com que as camadas criadas posteriormente usem precisão mista com uma mistura de float16 e float32.

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: Tesla V100-SXM2-16GB, compute capability 7.0

Resumindo, você pode passar diretamente uma string para set_global_policy , o que normalmente é feito na prática.

# Equivalent to the two lines above
mixed_precision.set_global_policy('mixed_float16')

A política especifica dois aspectos importantes de uma camada: o dtype em que os cálculos da camada são feitos e o dtype das variáveis ​​de uma camada. Acima, você criou uma política mixed_float16 (ou seja, uma mixed_precision.Policy criada passando a string 'mixed_float16' para seu construtor). Com esta política, as camadas usam cálculos float16 e variáveis ​​float32. Os cálculos são feitos em float16 para desempenho, mas as variáveis ​​devem ser mantidas em float32 para estabilidade numérica. Você pode consultar diretamente essas propriedades da política.

print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)

Compute dtype: float16
Variable dtype: float32

Como mencionado anteriormente, a política mixed_float16 melhorará significativamente o desempenho em GPUs NVIDIA com capacidade de computação de pelo menos 7.0. A política será executada em outras GPUs e CPUs, mas pode não melhorar o desempenho. Para TPUs, a política mixed_bfloat16 deve ser usada.

Construindo o modelo

Em seguida, vamos começar a construir um modelo simples. Modelos de brinquedos muito pequenos geralmente não se beneficiam da precisão mista, porque a sobrecarga do tempo de execução do TensorFlow geralmente domina o tempo de execução, tornando insignificante qualquer melhoria de desempenho na GPU. Portanto, vamos construir duas grandes camadas Dense com 4096 unidades cada se uma GPU for usada.

inputs = keras.Input(shape=(784,), name='digits')
if tf.config.list_physical_devices('GPU'):
  print('The model will run with 4096 units on a GPU')
  num_units = 4096
else:
  # Use fewer units on CPUs so the model finishes in a reasonable amount of time
  print('The model will run with 64 units on a CPU')
  num_units = 64
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)
dense2 = layers.Dense(num_units, activation='relu', name='dense_2')
x = dense2(x)

The model will run with 4096 units on a GPU

Cada camada tem uma política e usa a política global por padrão. Cada uma das camadas Dense , portanto, tem a política mixed_float16 porque você definiu a política global como mixed_float16 anteriormente. Isso fará com que as camadas densas façam cálculos float16 e tenham variáveis ​​float32. Eles lançam suas entradas para float16 para fazer cálculos float16, o que faz com que suas saídas sejam float16 como resultado. Suas variáveis ​​são float32 e serão convertidas em float16 quando as camadas forem chamadas para evitar erros de incompatibilidade de dtype.

print(dense1.dtype_policy)
print('x.dtype: %s' % x.dtype.name)
# 'kernel' is dense1's variable
print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)

<Policy "mixed_float16">
x.dtype: float16
dense1.kernel.dtype: float32

Em seguida, crie as previsões de saída. Normalmente, você pode criar as previsões de saída da seguinte maneira, mas isso nem sempre é numericamente estável com float16.

# INCORRECT: softmax and model output will be float16, when it should be float32
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)

Outputs dtype: float16

Uma ativação softmax no final do modelo deve ser float32. Como a política dtype é mixed_float16 , a ativação do softmax normalmente teria um dtype de computação float16 e tensores float16 de saída.

Isso pode ser corrigido separando as camadas Dense e softmax e passando dtype='float32' para a camada softmax:

# CORRECT: softmax and model output are float32
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)

Outputs dtype: float32

Passar dtype='float32' para o construtor de camada softmax substitui a política dtype da camada para ser a política float32 , que faz cálculos e mantém variáveis ​​em float32. De forma equivalente, você poderia ter passado dtype=mixed_precision.Policy('float32') ; camadas sempre convertem o argumento dtype em uma política. Como a camada de Activation não tem variáveis, a variável dtype da política é ignorada, mas o dtype de cálculo da política de float32 faz com que softmax e a saída do modelo sejam float32.

Adicionar um softmax float16 no meio de um modelo é bom, mas um softmax no final do modelo deve estar em float32. A razão é que se o tensor intermediário fluindo do softmax para a perda for float16 ou bfloat16, podem ocorrer problemas numéricos.

Você pode substituir o dtype de qualquer camada para ser float32 passando dtype='float32' se achar que não será numericamente estável com cálculos float16. Mas normalmente, isso só é necessário na última camada do modelo, pois a maioria das camadas tem precisão suficiente com mixed_float16 e mixed_bfloat16 .

Mesmo que o modelo não termine em um softmax, as saídas ainda devem ser float32. Embora desnecessário para este modelo específico, as saídas do modelo podem ser convertidas em float32 com o seguinte:

# The linear activation is an identity function. So this simply casts 'outputs'
# to float32. In this particular case, 'outputs' is already float32 so this is a
# no-op.
outputs = layers.Activation('linear', dtype='float32')(outputs)

Em seguida, termine e compile o modelo e gere os dados de entrada:

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255

Este exemplo converte os dados de entrada de int8 para float32. Você não converte para float16, pois a divisão por 255 está na CPU, que executa operações float16 mais lentas que as operações float32. Nesse caso, a diferença de desempenho é insignificante, mas em geral você deve executar a matemática de processamento de entrada em float32 se for executado na CPU. A primeira camada do modelo converterá as entradas para float16, pois cada camada converte entradas de ponto flutuante para seu tipo de computação.

Os pesos iniciais do modelo são recuperados. Isso permitirá treinar do zero novamente carregando os pesos.

initial_weights = model.get_weights()

Treinando o modelo com Model.fit

Em seguida, treine o modelo:

history = model.fit(x_train, y_train,
                    batch_size=8192,
                    epochs=5,
                    validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
print('Test loss:', test_scores[0])
print('Test accuracy:', test_scores[1])

Epoch 1/5
6/6 [==============================] - 2s 78ms/step - loss: 4.9609 - accuracy: 0.4132 - val_loss: 0.6643 - val_accuracy: 0.8437
Epoch 2/5
6/6 [==============================] - 0s 34ms/step - loss: 0.7752 - accuracy: 0.7789 - val_loss: 0.3098 - val_accuracy: 0.9175
Epoch 3/5
6/6 [==============================] - 0s 34ms/step - loss: 0.3620 - accuracy: 0.8848 - val_loss: 0.3149 - val_accuracy: 0.8969
Epoch 4/5
6/6 [==============================] - 0s 34ms/step - loss: 0.2998 - accuracy: 0.9066 - val_loss: 0.2988 - val_accuracy: 0.9068
Epoch 5/5
6/6 [==============================] - 0s 33ms/step - loss: 0.2298 - accuracy: 0.9285 - val_loss: 0.5062 - val_accuracy: 0.8414
313/313 - 0s - loss: 0.5163 - accuracy: 0.8392
Test loss: 0.5163048505783081
Test accuracy: 0.8392000198364258

Observe que o modelo imprime o tempo por etapa nos logs: por exemplo, "25ms/step". A primeira época pode ser mais lenta, pois o TensorFlow gasta algum tempo otimizando o modelo, mas depois o tempo por etapa deve se estabilizar.

Se você estiver executando este guia no Colab, poderá comparar o desempenho da precisão mista com float32. Para fazer isso, altere a política de mixed_float16 para float32 na seção "Definindo a política dtype" e execute novamente todas as células até este ponto. Em GPUs com capacidade de computação 7.X, você deve ver o tempo por etapa aumentar significativamente, indicando que a precisão mista acelerou o modelo. Certifique-se de alterar a política de volta para mixed_float16 e execute novamente as células antes de continuar com o guia.

Em GPUs com capacidade de computação de pelo menos 8.0 (GPUs Ampere e superiores), você provavelmente não verá nenhuma melhoria de desempenho no modelo de brinquedo neste guia ao usar precisão mista em comparação com float32. Isso se deve ao uso de TensorFloat-32 , que usa automaticamente matemática de menor precisão em certas operações float32 como tf.linalg.matmul . O TensorFloat-32 oferece algumas das vantagens de desempenho da precisão mista ao usar o float32. No entanto, em modelos do mundo real, você ainda verá melhorias significativas de desempenho de precisão mista devido à economia de largura de banda de memória e operações que o TensorFloat-32 não oferece suporte.

Se estiver executando a precisão mista em uma TPU, você não verá tanto ganho de desempenho em comparação com a execução da precisão mista em GPUs, especialmente GPUs pré-Ampere. Isso ocorre porque as TPUs fazem certas operações em bfloat16 sob o capô, mesmo com a política dtype padrão de float32. Isso é semelhante a como as GPUs Ampere usam o TensorFloat-32 por padrão. Em comparação com as GPUs Ampere, as TPUs normalmente obtêm menos ganhos de desempenho com precisão mista em modelos do mundo real.

Para muitos modelos do mundo real, a precisão mista também permite dobrar o tamanho do lote sem ficar sem memória, pois os tensores float16 ocupam metade da memória. No entanto, isso não se aplica a este modelo de brinquedo, pois você provavelmente pode executar o modelo em qualquer dtype em que cada lote consiste em todo o conjunto de dados MNIST de 60.000 imagens.

Escala de perda

A escala de perda é uma técnica que o tf.keras.Model.fit executa automaticamente com a política mixed_float16 para evitar underflow numérico. Esta seção descreve o que é dimensionamento de perda e a próxima seção descreve como usá-lo com um loop de treinamento personalizado.

Underflow e Overflow

O tipo de dados float16 tem uma faixa dinâmica estreita em comparação com float32. Isso significa que valores acima \(65504\) irão transbordar para o infinito e valores abaixo \(6.0 \times 10^{-8}\) irão para zero. float32 e bfloat16 têm uma faixa dinâmica muito maior para que overflow e underflow não sejam um problema.

Por exemplo:

x = tf.constant(256, dtype='float16')
(x ** 2).numpy()  # Overflow

inf

x = tf.constant(1e-5, dtype='float16')
(x ** 2).numpy()  # Underflow

0.0

Na prática, o estouro com float16 raramente ocorre. Além disso, o underflow também raramente ocorre durante o passe para frente. No entanto, durante a passagem para trás, os gradientes podem chegar a zero. O dimensionamento de perdas é uma técnica para evitar esse underflow.

Visão geral do dimensionamento de perdas

O conceito básico de escala de perda é simples: basta multiplicar a perda por algum número grande, digamos \(1024\), e você obtém o valor da escala de perda . Isso fará com que os gradientes também sejam dimensionados por \(1024\) , reduzindo bastante a chance de underflow. Uma vez que os gradientes finais são calculados, divida-os por \(1024\) para trazê-los de volta aos seus valores corretos.

O pseudocódigo para este processo é:

loss_scale = 1024
loss = model(inputs)
loss *= loss_scale
# Assume `grads` are float32. You do not want to divide float16 gradients.
grads = compute_gradient(loss, model.trainable_variables)
grads /= loss_scale

Escolher uma escala de perda pode ser complicado. Se a escala de perda for muito baixa, os gradientes ainda podem chegar a zero. Se for muito alto, ocorre o inverso do problema: os gradientes podem transbordar até o infinito.

Para resolver isso, o TensorFlow determina dinamicamente a escala de perda para que você não precise escolher uma manualmente. Se você usar tf.keras.Model.fit , o dimensionamento de perda é feito para você, então você não precisa fazer nenhum trabalho extra. Se você usar um loop de treinamento personalizado, deverá usar explicitamente o wrapper do otimizador especial tf.keras.mixed_precision.LossScaleOptimizer para usar o dimensionamento de perda. Isso é descrito na próxima seção.

Treinando o modelo com um loop de treinamento personalizado

Até agora, você treinou um modelo Keras com precisão mista usando tf.keras.Model.fit . Em seguida, você usará precisão mista com um loop de treinamento personalizado. Se você ainda não sabe o que é um loop de treinamento personalizado, leia primeiro o guia de treinamento personalizado .

Executar um loop de treinamento personalizado com precisão mista requer duas alterações ao executá-lo em float32:

  1. Construa o modelo com precisão mista (você já fez isso)
  2. Use explicitamente a escala de perda se mixed_float16 for usado.

Para a etapa (2), você usará a classe tf.keras.mixed_precision.LossScaleOptimizer , que envolve um otimizador e aplica a escala de perda. Por padrão, ele determina dinamicamente a escala de perda para que você não precise escolher uma. Construa um LossScaleOptimizer da seguinte maneira.

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)

Se desejar, é possível escolher uma escala de perda explícita ou personalizar o comportamento de escala de perda, mas é altamente recomendável manter o comportamento de escala de perda padrão, pois ele funciona bem em todos os modelos conhecidos. Consulte a documentação tf.keras.mixed_precision.LossScaleOptimizer se desejar personalizar o comportamento de dimensionamento de perda.

Em seguida, defina o objeto de perda e os tf.data.Dataset s:

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
                 .shuffle(10000).batch(8192))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)

Em seguida, defina a função da etapa de treinamento. Você usará dois novos métodos do otimizador de escala de perda para dimensionar a perda e reduzir a escala dos gradientes:

  • get_scaled_loss(loss) : Multiplica a perda pela escala de perda
  • get_unscaled_gradients(gradients) : Recebe uma lista de gradientes dimensionados como entradas e divide cada um pela escala de perda para desescalá-los

Estas funções devem ser utilizadas para evitar underflow nos gradientes. LossScaleOptimizer.apply_gradients então aplicará gradientes se nenhum deles tiver Inf s ou NaN s. Ele também atualizará a escala de perda, reduzindo-a pela metade se os gradientes tiverem Inf s ou NaN s e potencialmente aumentando-a caso contrário.

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

O LossScaleOptimizer provavelmente pulará as primeiras etapas no início do treinamento. A escala de perda começa alta para que a escala de perda ideal possa ser determinada rapidamente. Após algumas etapas, a escala de perda se estabilizará e poucas etapas serão ignoradas. Esse processo acontece automaticamente e não afeta a qualidade do treinamento.

Agora, defina a etapa de teste:

@tf.function
def test_step(x):
  return model(x, training=False)

Carregue os pesos iniciais do modelo, para que você possa treinar novamente do zero:

model.set_weights(initial_weights)

Por fim, execute o loop de treinamento personalizado:

for epoch in range(5):
  epoch_loss_avg = tf.keras.metrics.Mean()
  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='test_accuracy')
  for x, y in train_dataset:
    loss = train_step(x, y)
    epoch_loss_avg(loss)
  for x, y in test_dataset:
    predictions = test_step(x)
    test_accuracy.update_state(y, predictions)
  print('Epoch {}: loss={}, test accuracy={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))

Epoch 0: loss=4.869325160980225, test accuracy=0.7221999764442444
Epoch 1: loss=0.4893573224544525, test accuracy=0.878000020980835
Epoch 2: loss=0.36011582612991333, test accuracy=0.9440000057220459
Epoch 3: loss=0.27391332387924194, test accuracy=0.9318000078201294
Epoch 4: loss=0.247697651386261, test accuracy=0.933899998664856

Dicas de desempenho da GPU

Aqui estão algumas dicas de desempenho ao usar precisão mista em GPUs.

Aumentando o tamanho do seu lote

Se isso não afetar a qualidade do modelo, tente executar com o dobro do tamanho do lote ao usar a precisão mista. Como os tensores float16 usam metade da memória, isso geralmente permite dobrar o tamanho do lote sem ficar sem memória. Aumentar o tamanho do lote normalmente aumenta o rendimento do treinamento, ou seja, os elementos de treinamento por segundo em que seu modelo pode ser executado.

Garantir que os núcleos tensores da GPU sejam usados

Como mencionado anteriormente, as GPUs NVIDIA modernas usam uma unidade de hardware especial chamada Tensor Cores que pode multiplicar matrizes float16 muito rapidamente. No entanto, os Tensor Cores requerem que certas dimensões dos tensores sejam um múltiplo de 8. Nos exemplos abaixo, um argumento está em negrito se e somente se ele precisa ser um múltiplo de 8 para que os Tensor Cores sejam usados.

  • tf.keras.layers.Dense( unidades=64 )
  • tf.keras.layers.Conv2d( filtros=48 , kernel_size=7, stride=3)
    • E da mesma forma para outras camadas convolucionais, como tf.keras.layers.Conv3d
  • tf.keras.layers.LSTM( unidades=64 )
    • E semelhante para outras RNNs, como tf.keras.layers.GRU
  • tf.keras.Model.fit(epochs=2, batch_size=128 )

Você deve tentar usar Tensor Cores quando possível. Se você quiser saber mais, o guia de desempenho de aprendizado profundo da NVIDIA descreve os requisitos exatos para usar os Tensor Cores, bem como outras informações de desempenho relacionadas ao Tensor Core.

XLA

O XLA é um compilador que pode aumentar ainda mais o desempenho de precisão mista, bem como o desempenho do float32 em menor grau. Consulte o guia XLA para obter detalhes.

Dicas de desempenho do Cloud TPU

Assim como nas GPUs, você deve tentar dobrar o tamanho do lote ao usar Cloud TPUs porque os tensores bfloat16 usam metade da memória. Dobrar o tamanho do lote pode aumentar o rendimento do treinamento.

As TPUs não exigem nenhum outro ajuste específico de precisão mista para obter o desempenho ideal. Eles já exigem o uso de XLA. As TPUs se beneficiam de ter certas dimensões sendo múltiplos de \(128\), mas isso se aplica tanto ao tipo float32 quanto à precisão mista. Consulte o guia de desempenho do Cloud TPU para obter dicas gerais de desempenho do TPU, que se aplicam à precisão mista, bem como aos tensores float32.

Resumo

  • Você deve usar precisão mista se usar TPUs ou GPUs NVIDIA com pelo menos capacidade de computação 7.0, pois isso melhorará o desempenho em até 3x.

  • Você pode usar a precisão mista com as seguintes linhas:

    # On TPUs, use 'mixed_bfloat16' instead
mixed_precision.set_global_policy('mixed_float16')

  • Se o seu modelo terminar em softmax, certifique-se de que seja float32. E independentemente de como seu modelo termina, certifique-se de que a saída seja float32.

  • Se você usa um loop de treinamento personalizado com mixed_float16 , além das linhas acima, você precisa envolver seu otimizador com um tf.keras.mixed_precision.LossScaleOptimizer . Em seguida, chame optimizer.get_scaled_loss para dimensionar a perda e optimizer.get_unscaled_gradients para reduzir a escala dos gradientes.

  • Dobre o tamanho do lote de treinamento se isso não reduzir a precisão da avaliação

  • Em GPUs, certifique-se de que a maioria das dimensões do tensor seja um múltiplo de \(8\) para maximizar o desempenho

Para obter mais exemplos de precisão mista usando a API tf.keras.mixed_precision , consulte o repositório oficial de modelos . A maioria dos modelos oficiais, como ResNet e Transformer , serão executados usando precisão mista passando --dtype=fp16 .