mixed_precision.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Precisão mista

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

A precisão mista é o uso de tipos de ponto flutuante de 16 e 32 bits num modelo durante o treinamento para fazê-lo funcionar 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 passo menor e treinará igualmente em termos de métricas de avaliação, como exatidão. Este guia descreve como usar a API de precisão mista do Keras para acelerar seus modelos. O uso desta API pode melhorar o desempenho em mais de 3 vezes em GPUs modernas, 60% em TPUs e mais de 2 vezes nas CPUs Intel mais recentes.

Hoje, a maioria dos modelos usa o dtype float32, que ocupa 32 bits de memória. No entanto, existem dois dtypes de menor precisão, float16 e bfloat16, cada um ocupando 16 bits de memória. Os aceleradores modernos podem executar operações mais rapidamente nos dtypes de 16 bits, pois possuem hardware especializado para executar computações 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 rapidamente do que em float32, e TPUs e CPUs Intel suportados podem executar operações em bfloat16 mais rapidamente do em que float32. Portanto, esses dtypes de menor precisão devem ser usados ​​sempre que possível nesses dispositivos. No entanto, variáveis ​​e algumas computações ainda devem estar em float32 por motivos numéricos 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.

Observação: Neste guia, o termo "estabilidade numérica" ​​refere-se a como a qualidade de um modelo é afetada pelo uso de um dtype de menor precisão em vez de um dtype de maior precisão. Uma operação é "numericamente instável" em float16 ou bfloat16 se executá-la num desses dtypes fizer com que o modelo tenha pior exatidão de avaliação ou outras métricas em comparação com a execução da operação em float32.

Configuração

import tensorflow as tf

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

Hardware suportado

Embora a precisão mista funcione na maioria dos hardwares, ela só acelerará os modelos em GPUs NVIDIA recentes, TPUs de nuvem e CPUs Intel recentes. As GPUs NVIDIA suportam uma combinação de float16 e float32, enquanto TPUs e CPUs Intel suportam uma combinação de bfloat16 e float32.

Entre as GPUs NVIDIA, aquelas com capacidade de computação 7.0 ou superior apresentarão a maior vantagem de desempenho com a precisão mista porque possuem unidades de hardware especiais, chamadas Tensor Cores, para acelerar multiplicações e convoluções de matrizes float16. GPUs mais antigas não oferecem nenhuma vantagem de desempenho matemático para o uso de precisão mista; no entanto, a economia de memória e largura de banda pode permitir algumas acelerações. Você pode pesquisar a capacidade de computação da sua GPU na página CUDA GPU da NVIDIA. Exemplos de GPUs que mais se beneficiarão da precisão mista incluem GPUs RTX, V100 e A100.

As CPUs da Intel, começando com os processadores Intel Xeon de 4ª geração (codinome Sapphire Rapids), apresentarão o maior benefício de desempenho com precisão mista, pois podem acelerar cálculos bfloat16 usando instruções AMX (requer Tensorflow 2.12 ou posterior).

Observação: ao executar este guia no Google Colab, o runtime da GPU normalmente tem um P100 conectado. O P100 tem capacidade computacional 6.0 e não se espera que apresente uma aceleração significativa. Se estiver executando no runtime da CPU, pode haver lentidão, pois o runtime provavelmente tem uma CPU sem AMX.

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

!nvidia-smi -L

Todas as TPUs de nuvem suportam bfloat16.

Mesmo em CPUs Intel mais antigas, em outras CPUs x86 sem AMX e em GPUs mais antigas, onde nenhuma aceleração é esperada, APIs de precisão mista ainda podem ser usadas para testes de unidade, depuração ou apenas para testar a API. No entanto, mixed_bfloat16 em CPUs sem instruções AMX e mixed_float16 em todas as CPUs x86 serão executados significativamente mais devagar.

Configurando a política de dtypes

Para usar precisão mista em Keras, você precisa criar uma tf.keras.mixed_precision.Policy, normalmente chamada de dtype policy. As políticas de dtype especificam os dtypes nos quais as camadas serão executadas. Neste guia, você construirá uma política a partir da string 'mixed_float16' e a definirá como a política global. Isto 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)

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 computações float16 e variáveis ​​float32. As computações são feitas em float16 por questões de 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)

Conforme 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 poderá não melhorar o desempenho. Para TPUs e CPUs, a política mixed_bfloat16 deve ser usada.

Construindo o modelo

Agora vamos começar a construir um modelo simples. Modelos de brinquedo muito pequenos normalmente não se beneficiam da precisão mista, porque a sobrecarga do runtime do TensorFlow normalmente domina o tempo de execução, tornando qualquer melhoria de desempenho na GPU insignificante. Portanto, vamos construir duas grandes camadas Dense com 4.096 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)

Toda camada tem uma política e usará 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. Isto fará com que as camadas densas façam computações float16 e tenham variáveis ​​float32. Elas convertem suas entradas em float16 para fazer computações em 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)

Em seguida, crie as previsões de saída. Normalmente, você pode criar as previsões de saída da seguinte maneira, mas isto 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)

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

Isto 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)

Passar dtype='float32' para o construtor da camada softmax substitui a política dtype da camada pela política float32, que faz computações e mantém as 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 Activation não tem variáveis, a variável dtype da política é ignorada, mas a computação do dtype de float32 da política faz com que o softmax e a saída do modelo sejam float32.

Acrescentar 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 que flui do softmax para a perda for float16 ou bfloat16, poderão ocorrer problemas numéricos.

Você pode sobrepor o dtype de qualquer camada para float32 passando dtype='float32' se achar que não ele será numericamente estável com computações float16. Mas normalmente, isto só é necessário na última camada do modelo, já que a maioria das camadas tem precisão suficiente com mixed_float16 e mixed_bfloat16.

Mesmo que o modelo não termine em 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, finalize 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 faz o cast dos dados de entrada de int8 para float32. Você não faz cast para float16, pois a divisão por 255 está na CPU, que executa operações float16 mais lentamente que operações float32. Nesse caso, a diferença de desempenho é insignificante, mas em geral você deve executar a matemática do processamento de entrada em float32 se for executada na CPU. A primeira camada do modelo converterá as entradas para float16, à medida que cada camada converte entradas de ponto flutuante para seu dtype de computação.

Os pesos iniciais do modelo são recuperados. Isto 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])

Observe que o modelo imprime o tempo por cada passo nos logs: por exemplo, "25ms/step". A primeira época pode ser mais lenta, pois o TensorFlow passa algum tempo otimizando o modelo, mas depois o tempo por passo 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 "Configurando a política dtype" e execute novamente todas as células até este ponto. Em GPUs com capacidade de computação 7.X, você verá o tempo por passo aumentar significativamente, indicando que a precisão mista acelerou o modelo. Não esqueça de alterar a política de volta para mixed_float16 e executar 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 deste guia ao usar precisão mista em comparação com float32. Isto 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 float32. No entanto, em modelos do mundo real, você ainda experimentará melhorias significativas de desempenho com precisão mista devido à economia de largura de banda de memória e operações que o TensorFloat-32 não suporta.

Ao executar precisão mista numa TPU, você não verá tanto ganho de desempenho em comparação com a execução de precisão mista em GPUs, especialmente GPUs pré-Ampere. Isto ocorre porque as TPUs realizam certas operações em bfloat16 nos bastidores, 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 apresentam 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, já que os tensores float16 ocupam metade da memória. No entanto, isto não se aplica a este modelo de brinquedo, pois provavelmente você poderá executar o modelo em qualquer dtype em que cada lote consista em todo o dataset MNIST de 60.000 imagens.

Escalonamento de perdas

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

Observação: Ao usar a política mixed_bfloat16, não há necessidade de escalonar perdas.

Underflow e overflow

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

Por exemplo:

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

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

Na prática, o overflow com float16 raramente ocorre. Além disso, o underflow também raramente ocorre durante o passo para frente. No entanto, durante o passo para trás, os gradientes podem sofrer um underflow para zero. O escalonamento de perdas é uma técnica para evitar esse underflow.

Visão geral do escalonamento de perdas

O conceito básico do escalonamento de perdas é simples: basta multiplicar a perda por algum número grande, digamos $1024$, que você obterá o valor da escala de perdas. Isto fará com que os gradientes também aumentem em $1024$, reduzindo bastante a chance de underflow. Depois que os gradientes finais forem calculados, divida-os por $1024$ para trazê-los de volta aos 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 perdas pode ser complicado. Se a escala de perdas for muito baixa, poderá ainda ocorrer um underflow dos gradientes a zero. Se for muito alto, ocorre o problema oposto: poderá haver overflow dos gradientes até o infinito.

Para solucionar esse problema, 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 escalonamento de perdas será feito para você, para que você não precise fazer nenhum trabalho extra. Se você usar um loop de treinamento personalizado, deverá usar explicitamente o wrapper especial do otimizador tf.keras.mixed_precision.LossScaleOptimizer para usar o escalonamento de perdas. 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. A seguir, você usará precisão mista com um loop de treinamento personalizado. Se você ainda não sabe o que é um ciclo de treinamento personalizado, leia primeiro o guia de treinamento personalizado.

A execução de um loop de treinamento personalizado com precisão mista requer duas alterações em relação à execução em float32:

1. Construir o modelo com precisão mista (você já fez isso)

2. Usar o escalonamento de perdas explicitamente se mixed_float16 for usado.

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

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

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

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

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)

Depois defina a função do passo de treinamento. Você usará dois novos métodos do otimizador de escala de perdas para escalonar a perda e desescalonar os gradientes:

• get_scaled_loss(loss): multiplica a perda pela escala de perda

• get_unscaled_gradients(gradients): recebe uma lista de gradientes escalonados como entradas e divide cada um pela escala de perda para desescaloná-los

Essas funções devem ser utilizadas para evitar underflow nos gradientes. LossScaleOptimizer.apply_gradients aplicará gradientes se nenhum deles tiver valores Inf ou NaN. Ele também atualizará a escala de perdas, reduzindo-a pela metade se os gradientes tiverem valores Inf ou NaN 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á os primeiros passos no início do treinamento. A escala de perdas começa alta para que a escala ideal possa ser determinada rapidamente. Após alguns passos, a escala de perdas se estabilizará e poucos passos serão ignorados. Este processo acontece automaticamente e não afeta a qualidade do treinamento.

Agora, defina o passo de teste:

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

Carregue os pesos iniciais do modelo para poder 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()))

Dicas de desempenho da GPU

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

Aumentando o tamanho do lote

Se isso não afetar a qualidade do modelo, tente executar com o dobro do tamanho do lote ao usar precisão mista. Como os tensores float16 usam metade da memória, isto 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.

Garantindo que os Tensor Cores da GPU sejam usados

Conforme 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 exigem que certas dimensões dos tensores sejam múltiplos de 8. Nos exemplos abaixo, um argumento fica em negrito se e somente se precisar ser um múltiplo de 8 para que os Tensor Cores sejam usados.

• tf.keras.layers.Dense(units=64)

• tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)

  • E da mesma forma para outras camadas convolucionais, como tf.keras.layers.Conv3d

• tf.keras.layers.LSTM(units=64)

  • E da mesma forma para outras RNNs, como tf.keras.layers.GRU

• tf.keras.Model.fit(epochs=2, batch_size=128)

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

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 mais detalhes.

Dicas de desempenho do TPU na nuvem

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

As TPUs não requerem nenhum outro ajuste específico de precisão mista para obter o desempenho ideal. Elas já exigem o uso de XLA. As TPUs se beneficiam do fato de certas dimensões serem múltiplos de $128$, mas isto se aplica igualmente ao tipo float32 e à precisão mista. Consulte o guia de desempenho da TPU na nuvem para dicas gerais de desempenho para TPUs, que se aplicam à precisão mista e também aos tensores float32.

Resumo

• Você deve usar precisão mista se usar TPUs, GPUs NVIDIA com capacidade de computação pelo menos 7.0 ou CPUs Intel com suporte para instruções AMX, pois isto melhorará o desempenho em até três vezes.

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

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

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

• Se você usar um loop de treinamento personalizado com mixed_float16, além das linhas acima, será necessário empacotar seu otimizador com um wrapper tf.keras.mixed_precision.LossScaleOptimizer. Em seguida, chame optimizer.get_scaled_loss para escalonar a perda e optimizer.get_unscaled_gradients para desescalonar os gradientes.

• Se você usa um loop de treinamento personalizado com mixed_bfloat16, definir a global_policy mencionada acima será suficiente.

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

• Em GPUs, certifique-se de que a maioria das dimensões do tensor sejam múltiplos de $8$ para maximizar o desempenho

Para um exemplo de precisão mista usando a API tf.keras.mixed_precision, veja funções e classes relacionadas ao desempenho do treinamento. Confira os modelos oficiais, como o Transformer, para mais detalhes.