Como escrever loops de treinamento personalizados em keras com gradientetape

Visão geral do conteúdo

• Configurar
• Introdução
• Usando o Gradienttape: um primeiro exemplo de ponta a ponta
• Manuseio de métricas de baixo nível
• Acelerando sua etapa de treinamento com tf.função
• Manuseio de baixo nível de perdas rastreadas pelo modelo
• Resumo
• Exemplo de ponta a ponta: um loop de treinamento GaN do zero

Configurar

import tensorflow as tf
import keras
from keras import layers
import numpy as np

Introdução

Keras fornece loops de treinamento e avaliação padrão, fit() e evaluate(). O uso deles é abordado no treinamento e avaliação do guia com os métodos internos.

Se você deseja personalizar o algoritmo de aprendizado do seu modelo, ainda aproveitando a conveniência de fit() (por exemplo, para treinar um gan usando fit()), você pode subclasse o Model classe e implemente o seu próprio train_step() método, que é chamado repetidamente durante fit(). Isso é abordado no guia personalizando o que acontece em fit().

Agora, se você deseja um controle de nível muito baixo sobre o treinamento e a avaliação, escreva seus próprios loops de treinamento e avaliação do zero. É disso que se trata este guia.

Usando o GradientTape: Um exemplo de ponta a ponta

Chamando um modelo dentro de um GradientTape O escopo permite recuperar os gradientes dos pesos treináveis ​​da camada em relação a um valor de perda. Usando uma instância do otimizador, você pode usar esses gradientes para atualizar essas variáveis ​​(que você pode recuperar usando model.trainable_weights).

Vamos considerar um modelo mnist simples:

inputs = keras.Input(shape=(784,), name="digits")
x1 = layers.Dense(64, activation="relu")(inputs)
x2 = layers.Dense(64, activation="relu")(x1)
outputs = layers.Dense(10, name="predictions")(x2)
model = keras.Model(inputs=inputs, outputs=outputs)

Vamos treiná-lo usando o mini gradiente em loop de treinamento personalizado.

Primeiro, precisaremos de um otimizador, uma função de perda e um conjunto de dados:

# Instantiate an optimizer.
optimizer = keras.optimizers.SGD(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Prepare the training dataset.
batch_size = 64
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784))
x_test = np.reshape(x_test, (-1, 784))

# Reserve 10,000 samples for validation.
x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]

# Prepare the training dataset.
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

# Prepare the validation dataset.
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(batch_size)

Aqui está o nosso loop de treinamento:

• Abrimos um for loop que itera sobre épocas
• Para cada época, abrimos um for loop que itera sobre o conjunto de dados, em lotes
• Para cada lote, abrimos um GradientTape() escopo
• Dentro deste escopo, chamamos o modelo (passagem para a frente) e calculamos a perda
• Fora do escopo, recuperamos os gradientes dos pesos do modelo em relação à perda
• Por fim, usamos o otimizador para atualizar os pesos do modelo com base nos gradientes

epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        # Open a GradientTape to record the operations run
        # during the forward pass, which enables auto-differentiation.
        with tf.GradientTape() as tape:
            # Run the forward pass of the layer.
            # The operations that the layer applies
            # to its inputs are going to be recorded
            # on the GradientTape.
            logits = model(x_batch_train, training=True)  # Logits for this minibatch

            # Compute the loss value for this minibatch.
            loss_value = loss_fn(y_batch_train, logits)

        # Use the gradient tape to automatically retrieve
        # the gradients of the trainable variables with respect to the loss.
        grads = tape.gradient(loss_value, model.trainable_weights)

        # Run one step of gradient descent by updating
        # the value of the variables to minimize the loss.
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %s samples" % ((step + 1) * batch_size))

Start of epoch 0
WARNING:tensorflow:5 out of the last 5 calls to  triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to  and  for  more details.
WARNING:tensorflow:6 out of the last 6 calls to  triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to  and  for  more details.
Training loss (for one batch) at step 0: 131.3794
Seen so far: 64 samples
Training loss (for one batch) at step 200: 1.2871
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 1.2652
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.8800
Seen so far: 38464 samples

Start of epoch 1
Training loss (for one batch) at step 0: 0.8296
Seen so far: 64 samples
Training loss (for one batch) at step 200: 1.3322
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 1.0486
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.6610
Seen so far: 38464 samples

Manuseio de métricas de baixo nível

Vamos adicionar o monitoramento de métricas a esse loop básico.

Você pode reutilizar prontamente as métricas embutidas (ou as personalizadas que você escreveu) em tais loops de treinamento escritos do zero. Aqui está o fluxo:

• Instanciar a métrica no início do loop
• Chamar metric.update_state() Após cada lote
• Chamar metric.result() Quando você precisa exibir o valor atual da métrica
• Chamar metric.reset_states() Quando você precisa limpar o estado da métrica (normalmente no final de uma época)

Vamos usar esse conhecimento para calcular SparseCategoricalAccuracy nos dados de validação no final de cada época:

# Get model
inputs = keras.Input(shape=(784,), name="digits")
x = layers.Dense(64, activation="relu", name="dense_1")(inputs)
x = layers.Dense(64, activation="relu", name="dense_2")(x)
outputs = layers.Dense(10, name="predictions")(x)
model = keras.Model(inputs=inputs, outputs=outputs)

# Instantiate an optimizer to train the model.
optimizer = keras.optimizers.SGD(learning_rate=1e-3)
# Instantiate a loss function.
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Prepare the metrics.
train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = keras.metrics.SparseCategoricalAccuracy()

Aqui está o nosso loop de treinamento e avaliação:

import time

epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            logits = model(x_batch_train, training=True)
            loss_value = loss_fn(y_batch_train, logits)
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

        # Update training metric.
        train_acc_metric.update_state(y_batch_train, logits)

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %d samples" % ((step + 1) * batch_size))

    # Display metrics at the end of each epoch.
    train_acc = train_acc_metric.result()
    print("Training acc over epoch: %.4f" % (float(train_acc),))

    # Reset training metrics at the end of each epoch
    train_acc_metric.reset_states()

    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:
        val_logits = model(x_batch_val, training=False)
        # Update val metrics
        val_acc_metric.update_state(y_batch_val, val_logits)
    val_acc = val_acc_metric.result()
    val_acc_metric.reset_states()
    print("Validation acc: %.4f" % (float(val_acc),))
    print("Time taken: %.2fs" % (time.time() - start_time))

Start of epoch 0
Training loss (for one batch) at step 0: 106.2691
Seen so far: 64 samples
Training loss (for one batch) at step 200: 0.9259
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 0.9347
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.7641
Seen so far: 38464 samples
Training acc over epoch: 0.7332
Validation acc: 0.8325
Time taken: 10.95s

Start of epoch 1
Training loss (for one batch) at step 0: 0.5238
Seen so far: 64 samples
Training loss (for one batch) at step 200: 0.7125
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 0.5705
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.6006
Seen so far: 38464 samples
Training acc over epoch: 0.8424
Validation acc: 0.8525
Time taken: 10.59s

Acelerando sua etapa de treinamento com tf.function

O tempo de execução padrão no Tensorflow 2 é a execução ansiosa. Como tal, nosso ciclo de treinamento acima é executado ansiosamente.

Isso é ótimo para depuração, mas a compilação de gráficos tem uma vantagem definitiva de desempenho. Descrever seu cálculo como um gráfico estático permite que a estrutura aplique otimizações globais de desempenho. Isso é impossível quando a estrutura é restrita a executar avidamente uma operação após a outra, sem nenhum conhecimento do que vem a seguir.

Você pode compilar em um gráfico estático qualquer função que tome tensores como entrada. Basta adicionar um @tf.function Decorador, assim:

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    train_acc_metric.update_state(y, logits)
    return loss_value

Vamos fazer o mesmo com a etapa de avaliação:

@tf.function
def test_step(x, y):
    val_logits = model(x, training=False)
    val_acc_metric.update_state(y, val_logits)

Agora, vamos executar novamente nosso loop de treinamento com esta etapa de treinamento compilada:

import time

epochs = 2
for epoch in range(epochs):
    print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_value = train_step(x_batch_train, y_batch_train)

        # Log every 200 batches.
        if step % 200 == 0:
            print(
                "Training loss (for one batch) at step %d: %.4f"
                % (step, float(loss_value))
            )
            print("Seen so far: %d samples" % ((step + 1) * batch_size))

    # Display metrics at the end of each epoch.
    train_acc = train_acc_metric.result()
    print("Training acc over epoch: %.4f" % (float(train_acc),))

    # Reset training metrics at the end of each epoch
    train_acc_metric.reset_states()

    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:
        test_step(x_batch_val, y_batch_val)

    val_acc = val_acc_metric.result()
    val_acc_metric.reset_states()
    print("Validation acc: %.4f" % (float(val_acc),))
    print("Time taken: %.2fs" % (time.time() - start_time))

Start of epoch 0
Training loss (for one batch) at step 0: 0.5162
Seen so far: 64 samples
Training loss (for one batch) at step 200: 0.4599
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 0.3975
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.2557
Seen so far: 38464 samples
Training acc over epoch: 0.8747
Validation acc: 0.8545
Time taken: 1.85s

Start of epoch 1
Training loss (for one batch) at step 0: 0.6145
Seen so far: 64 samples
Training loss (for one batch) at step 200: 0.3751
Seen so far: 12864 samples
Training loss (for one batch) at step 400: 0.3464
Seen so far: 25664 samples
Training loss (for one batch) at step 600: 0.4128
Seen so far: 38464 samples
Training acc over epoch: 0.8919
Validation acc: 0.8996
Time taken: 1.34s

Muito mais rápido, não é?

Manuseio de baixo nível de perdas rastreadas pelo modelo

Camadas e modelos rastreiam recursivamente quaisquer perdas criadas durante o passe para a frente por camadas que chamam self.add_loss(value). A lista resultante de valores de perda escalar estão disponíveis através da propriedade model.losses No final do passe para a frente.

Se você deseja usar esses componentes de perda, deve resumir e adicioná -los à principal perda na sua etapa de treinamento.

Considere esta camada, que cria uma perda de regularização da atividade:

@keras.saving.register_keras_serializable()
class ActivityRegularizationLayer(layers.Layer):
    def call(self, inputs):
        self.add_loss(1e-2 * tf.reduce_sum(inputs))
        return inputs

Vamos construir um modelo realmente simples que o use:

inputs = keras.Input(shape=(784,), name="digits")
x = layers.Dense(64, activation="relu")(inputs)
# Insert activity regularization as a layer
x = ActivityRegularizationLayer()(x)
x = layers.Dense(64, activation="relu")(x)
outputs = layers.Dense(10, name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs)

Aqui está a aparência da nossa etapa de treinamento agora:

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
        # Add any extra losses created during the forward pass.
        loss_value += sum(model.losses)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    train_acc_metric.update_state(y, logits)
    return loss_value

Resumo

Agora você sabe tudo o que há para saber sobre o uso de loops de treinamento interno e escrever o seu próprio do zero.

Para concluir, aqui está um exemplo simples de ponta a ponta que une tudo o que você aprendeu neste guia: um DCGAN treinado em dígitos mnist.

Exemplo de ponta a ponta: um loop de treinamento GaN do zero

Você pode estar familiarizado com redes adversárias generativas (GANS). Os Gans podem gerar novas imagens que parecem quase reais, aprendendo a distribuição latente de um conjunto de dados de treinamento de imagens (o “espaço latente” das imagens).

Um GaN é feito de duas partes: um modelo “gerador” que mapeia pontos no espaço latente para pontos no espaço da imagem, um modelo “discriminador”, um classificador que pode dizer a diferença entre imagens reais (do conjunto de dados de treinamento) e imagens falsas (a saída da rede geradora).

Um loop de treinamento GaN se parece com o seguinte:

1. Treine o discriminador. – Amostra um lote de pontos aleatórios no espaço latente. – Transforme os pontos em imagens falsas através do modelo “gerador”. – Obtenha um lote de imagens reais e combine -as com as imagens geradas. – Treine o modelo “discriminador” para classificar imagens geradas versus reais.
2. Treine o gerador. – Amostra pontos aleatórios no espaço latente. – Transforme os pontos em imagens falsas através da rede “gerador”. – Obtenha um lote de imagens reais e combine -as com as imagens geradas. – Treine o modelo “gerador” para “enganar” o discriminador e classificar as imagens falsas como reais.

Para uma visão geral muito mais detalhada de como os Gans funcionam, consulte Deep Learning With Python.

Vamos implementar esse loop de treinamento. Primeiro, crie o discriminador destinado a classificar dígitos falsos vs reais:

discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.GlobalMaxPooling2D(),
        layers.Dense(1),
    ],
    name="discriminator",
)
discriminator.summary()

Model: "discriminator"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d (Conv2D)             (None, 14, 14, 64)        640       
                                                                 
 leaky_re_lu (LeakyReLU)     (None, 14, 14, 64)        0         
                                                                 
 conv2d_1 (Conv2D)           (None, 7, 7, 128)         73856     
                                                                 
 leaky_re_lu_1 (LeakyReLU)   (None, 7, 7, 128)         0         
                                                                 
 global_max_pooling2d (Glob  (None, 128)               0         
 alMaxPooling2D)                                                 
                                                                 
 dense_4 (Dense)             (None, 1)                 129       
                                                                 
=================================================================
Total params: 74625 (291.50 KB)
Trainable params: 74625 (291.50 KB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________

Então vamos criar uma rede de geradores, que transforma vetores latentes em saídas de forma (28, 28, 1) (representando dígitos mnist):

latent_dim = 128

generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        # We want to generate 128 coefficients to reshape into a 7x7x128 map
        layers.Dense(7 * 7 * 128),
        layers.LeakyReLU(alpha=0.2),
        layers.Reshape((7, 7, 128)),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
    ],
    name="generator",
)

Aqui está a parte chave: o loop de treinamento. Como você pode ver, é bastante direto. A função da etapa de treinamento leva apenas 17 linhas.

# Instantiate one optimizer for the discriminator and another for the generator.
d_optimizer = keras.optimizers.Adam(learning_rate=0.0003)
g_optimizer = keras.optimizers.Adam(learning_rate=0.0004)

# Instantiate a loss function.
loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)


@tf.function
def train_step(real_images):
    # Sample random points in the latent space
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # Decode them to fake images
    generated_images = generator(random_latent_vectors)
    # Combine them with real images
    combined_images = tf.concat([generated_images, real_images], axis=0)

    # Assemble labels discriminating real from fake images
    labels = tf.concat(
        [tf.ones((batch_size, 1)), tf.zeros((real_images.shape[0], 1))], axis=0
    )
    # Add random noise to the labels - important trick!
    labels += 0.05 * tf.random.uniform(labels.shape)

    # Train the discriminator
    with tf.GradientTape() as tape:
        predictions = discriminator(combined_images)
        d_loss = loss_fn(labels, predictions)
    grads = tape.gradient(d_loss, discriminator.trainable_weights)
    d_optimizer.apply_gradients(zip(grads, discriminator.trainable_weights))

    # Sample random points in the latent space
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    # Assemble labels that say "all real images"
    misleading_labels = tf.zeros((batch_size, 1))

    # Train the generator (note that we should *not* update the weights
    # of the discriminator)!
    with tf.GradientTape() as tape:
        predictions = discriminator(generator(random_latent_vectors))
        g_loss = loss_fn(misleading_labels, predictions)
    grads = tape.gradient(g_loss, generator.trainable_weights)
    g_optimizer.apply_gradients(zip(grads, generator.trainable_weights))
    return d_loss, g_loss, generated_images

Vamos treinar nosso gan, ligando repetidamente train_step em lotes de imagens.

Como nosso discriminador e gerador são convênios, você vai querer executar esse código em uma GPU.

import os

# Prepare the dataset. We use both the training & test MNIST digits.
batch_size = 64
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
all_digits = np.concatenate([x_train, x_test])
all_digits = all_digits.astype("float32") / 255.0
all_digits = np.reshape(all_digits, (-1, 28, 28, 1))
dataset = tf.data.Dataset.from_tensor_slices(all_digits)
dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)

epochs = 1  # In practice you need at least 20 epochs to generate nice digits.
save_dir = "./"

for epoch in range(epochs):
    print("\nStart epoch", epoch)

    for step, real_images in enumerate(dataset):
        # Train the discriminator & generator on one batch of real images.
        d_loss, g_loss, generated_images = train_step(real_images)

        # Logging.
        if step % 200 == 0:
            # Print metrics
            print("discriminator loss at step %d: %.2f" % (step, d_loss))
            print("adversarial loss at step %d: %.2f" % (step, g_loss))

            # Save one generated image
            img = keras.utils.array_to_img(generated_images[0] * 255.0, scale=False)
            img.save(os.path.join(save_dir, "generated_img" + str(step) + ".png"))

        # To limit execution time we stop after 10 steps.
        # Remove the lines below to actually train the model!
        if step > 10:
            break

Start epoch 0
discriminator loss at step 0: 0.72
adversarial loss at step 0: 0.72

É isso! Você terá dígitos MNIST falsos bonitos após apenas ~ 30 anos de treinamento na GPU Colab.