Guia para escrever retornos de chamada personalizados do TensorFlow/Keras

Introdução

Suponha que você queira que seu modelo Keras tenha algum comportamento específico durante o treinamento, avaliação ou previsão. Por exemplo, você pode querer salvar seu modelo em cada época de treinamento. Uma maneira de fazer isso é usando Callbacks.

Em geral, Callbacks são funções que são chamadas quando algum evento acontece e são passadas como argumentos para outras funções. No caso do Keras, eles são uma ferramenta para customizar o comportamento do seu modelo – seja durante o treinamento, avaliação ou inferência. Alguns aplicativos estão registrando, persistindo no modelo, interrompendo antecipadamente ou alterando a taxa de aprendizado. Isso é feito passando uma lista de Callbacks como argumentos para keras.Model.fit(),keras.Model.evaluate() or keras.Model.predict().

Alguns casos de uso comuns para retornos de chamada estão modificando a taxa de aprendizado, registro, monitoramento e interrupção antecipada do treinamento. Keras tem uma série de retornos de chamada integrados, detalhados
na documentação.

No entanto, alguns aplicativos mais específicos podem exigir um retorno de chamada personalizado. Por exemplo, implementando o aquecimento da taxa de aprendizado com um decaimento do cosseno após um período de espera não está embutido no momento, mas é amplamente usado e adotado como um agendador.

Classe de retorno de chamada e seus métodos

Keras tem uma classe de retorno de chamada específica, keras.callbacks.Callback, com métodos que podem ser chamados durante o treinamento, teste e inferência em nível global, de lote ou de época. Em ordem de criar retornos de chamada personalizados, precisamos criar uma subclasse e substituir esses métodos.

A keras.callbacks.Callback classe tem três tipos de métodos:

• métodos globais: chamados no início ou no final de fit(), evaluate() e predict().
• métodos em nível de lote: chamados no início ou no final do processamento de um lote.
• métodos de nível de época: chamados no início ou no final de um lote de treinamento.

Observação: Cada método tem acesso a um dict chamado logs. As chaves e valores de logs são contextuais – eles dependem do evento que chama o método. Além disso, temos acesso ao modelo dentro de cada método através do self.model atributo.

Vamos dar uma olhada em três exemplos de callbacks personalizados – um para treinamento, um para avaliação e outro para previsão. Cada um imprimirá em cada etapa o que nosso modelo está fazendo e quais logs temos acesso. Isso é útil para entender o que é possível fazer com retornos de chamada personalizados em cada estágio.

Vamos começar definindo um modelo de brinquedo:

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

model = keras.Sequential()
model.add(keras.layers.Dense(10, input_dim = 1, activation='relu'))
model.add(keras.layers.Dense(10, activation='relu'))
model.add(keras.layers.Dense(1))
model.compile(
    optimizer=keras.optimizers.RMSprop(learning_rate=0.1),
    loss = "mean_squared_error",
    metrics = ["mean_absolute_error"]
)

x = np.random.uniform(low = 0, high = 10, size = 1000)
y = x**2
x_train, x_test = (x[:900],x[900:])
y_train, y_test = (y[:900],y[900:])

Retorno de chamada de treinamento personalizado

Nosso primeiro retorno de chamada deve ser chamado durante o treinamento. Vamos subclassificar o Callback classe:

class TrainingCallback(keras.callbacks.Callback):
    def __init__(self):
        self.tabulation = {"train":"", 'batch': " "*8, 'epoch':" "*4}
    def on_train_begin(self, logs=None):
        tab = self.tabulation['train']
        print(f"{tab}Training!")
        print(f"{tab}available logs: {logs}")

    def on_train_batch_begin(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}Batch {batch}")
        print(f"{tab}available logs: {logs}")

    def on_train_batch_end(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}End of Batch {batch}")
        print(f"{tab}available logs: {logs}")

    def on_epoch_begin(self, epoch, logs=None):
        tab = self.tabulation['epoch']
        print(f"{tab}Epoch {epoch} of training")
        print(f"{tab}available logs: {logs}")

    def on_epoch_end(self, epoch, logs=None):
        tab = self.tabulation['epoch']
        print(f"{tab}End of Epoch {epoch} of training")
        print(f"{tab}available logs: {logs}")

    def on_train_end(self, logs=None):
        tab = self.tabulation['train']
        print(f"{tab}Finishing training!")
        print(f"{tab}available logs: {logs}")

Se algum desses métodos não for substituído, o comportamento padrão continuará como antes. Em nosso exemplo – simplesmente imprimimos os logs disponíveis e o nível em que o retorno de chamada é aplicado, com o recuo adequado.

Vamos dar uma olhada nas saídas:

model.fit(
    x_train,
    y_train,
    batch_size=500,
    epochs=2,
    verbose=0,
    callbacks=[TrainingCallback()],
)

Training!
available logs: {}
    Epoch 0 of training
    available logs: {}
        Batch 0
        available logs: {}
        End of Batch 0
        available logs: {'loss': 2172.373291015625, 'mean_absolute_error': 34.79669952392578}
        Batch 1
        available logs: {}
        End of Batch 1
        available logs: {'loss': 2030.1309814453125, 'mean_absolute_error': 33.30256271362305}
    End of Epoch 0 of training
    available logs: {'loss': 2030.1309814453125, 'mean_absolute_error': 33.30256271362305}
    Epoch 1 of training
    available logs: {}
        Batch 0
        available logs: {}
        End of Batch 0
        available logs: {'loss': 1746.2772216796875, 'mean_absolute_error': 30.268001556396484}
        Batch 1
        available logs: {}
        End of Batch 1
        available logs: {'loss': 1467.36376953125, 'mean_absolute_error': 27.10252571105957}
    End of Epoch 1 of training
    available logs: {'loss': 1467.36376953125, 'mean_absolute_error': 27.10252571105957}
Finishing training!
available logs: {'loss': 1467.36376953125, 'mean_absolute_error': 27.10252571105957}

Observe que podemos acompanhar a cada etapa o que o modelo está fazendo e a quais métricas temos acesso. Ao final de cada lote e época, temos acesso à função de perda na amostra e às métricas do nosso modelo.

Retorno de chamada de avaliação personalizada

Agora, vamos chamar o Model.evaluate() método. Podemos ver que ao final de um lote temos acesso à função de perda e as métricas do momento, e ao final da avaliação temos acesso à perda e métricas gerais:

class TestingCallback(keras.callbacks.Callback):
    def __init__(self):
          self.tabulation = {"test":"", 'batch': " "*8}
      
    def on_test_begin(self, logs=None):
        tab = self.tabulation['test']
        print(f'{tab}Evaluating!')
        print(f'{tab}available logs: {logs}')

    def on_test_end(self, logs=None):
        tab = self.tabulation['test']
        print(f'{tab}Finishing evaluation!')
        print(f'{tab}available logs: {logs}')

    def on_test_batch_begin(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}Batch {batch}")
        print(f"{tab}available logs: {logs}")

    def on_test_batch_end(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}End of batch {batch}")
        print(f"{tab}available logs: {logs}")

res = model.evaluate(
    x_test, y_test, batch_size=100, verbose=0, callbacks=[TestingCallback()]
)

Evaluating!
available logs: {}
        Batch 0
        available logs: {}
        End of batch 0
        available logs: {'loss': 382.2723083496094, 'mean_absolute_error': 14.069927215576172}
Finishing evaluation!
available logs: {'loss': 382.2723083496094, 'mean_absolute_error': 14.069927215576172}

Retorno de chamada de previsão personalizado

Finalmente, vamos chamar o Model.predict() método. Observe que ao final de cada lote temos acesso às saídas previstas do nosso modelo:

class PredictionCallback(keras.callbacks.Callback):
    def __init__(self):
        self.tabulation = {"prediction":"", 'batch': " "*8}

    def on_predict_begin(self, logs=None):
        tab = self.tabulation['prediction']
        print(f"{tab}Predicting!")
        print(f"{tab}available logs: {logs}")

    def on_predict_end(self, logs=None):
        tab = self.tabulation['prediction']
        print(f"{tab}End of Prediction!")
        print(f"{tab}available logs: {logs}")

    def on_predict_batch_begin(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}batch {batch}")
        print(f"{tab}available logs: {logs}")

    def on_predict_batch_end(self, batch, logs=None):
        tab = self.tabulation['batch']
        print(f"{tab}End of batch {batch}")
        print(f"{tab}available logs:n {logs}")

res = model.predict(x_test[:10],
                    verbose = 0, 
                    callbacks=[PredictionCallback()])

Predicting!
available logs: {}
        batch 0
        available logs: {}
        End of batch 0
        available logs:
 {'outputs': array([[ 7.743822],
       [27.748264],
       [33.082104],
       [26.530678],
       [27.939169],
       [18.414223],
       [42.610645],
       [36.69335 ],
       [13.096557],
       [37.120853]], dtype=float32)}
End of Prediction!
available logs: {}

Com estes – você pode personalizar o comportamento, configurar o monitoramento ou alterar os processos de treinamento, avaliação ou inferência. Uma alternativa à subclassificação é usar o LambdaCallback.

Usando LambaCallback

Um dos retornos de chamada embutidos no Keras é o LambdaCallback classe. Este retorno de chamada aceita uma função que define como ela se comporta e o que ela faz! De certa forma, ele permite que você use qualquer função arbitrária como retorno de chamada, permitindo assim criar retornos de chamada personalizados.

A classe tem os parâmetros opcionais:
-on_epoch_begin

• on_epoch_end
• on_batch_begin
• on_batch_end
• on_train_begin
• on_train_end

Cada parâmetro aceita uma função que é chamado no respectivo evento de modelo. Como exemplo, vamos fazer um callback para enviar um email quando o modelo terminar o treinamento:

import smtplib
from email.message import EmailMessage

def send_email(logs): 
    msg = EmailMessage()
    content = f"""The model has finished training."""
    for key, value in logs.items():
      content = content + f"n{key}:{value:.2f}"
    msg.set_content(content)
    msg['Subject'] = f'Training report'
    msg['From'] = '[email protected]'
    msg['To'] = 'receiver-email'

    s = smtplib.SMTP('smtp.gmail.com', 587)
    s.starttls()
    s.login("[email protected]", "your-gmail-app-password")
    s.send_message(msg)
    s.quit()

lambda_send_email = lambda logs : send_email(logs)

email_callback = keras.callbacks.LambdaCallback(on_train_end = lambda_send_email)

model.fit(
    x_train,
    y_train,
    batch_size=100,
    epochs=1,
    verbose=0,
    callbacks=[email_callback],
)

Para fazer nosso retorno de chamada personalizado usando LambdaCallback, precisamos apenas implementar a função que queremos ser chamada, envolvê-la como um lambda função e passá-lo para o
LambdaCallback classe como um parâmetro.

Um retorno de chamada para visualizar o treinamento do modelo

Nesta seção, daremos um exemplo de um retorno de chamada personalizado que faz uma animação do desempenho do nosso modelo melhorando durante o treinamento. Para isso, armazenamos os valores dos logs no final de cada lote. Então, no final do loop de treinamento, criamos uma animação usando matplotlib.

Para melhorar a visualização, a perda e as métricas serão plotadas em escala logarítmica:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from IPython import display

class TrainingAnimationCallback(keras.callbacks.Callback):
    def __init__(self, duration = 40, fps = 1000/25):
        self.duration = duration
        self.fps = fps
        self.logs_history = []

    def set_plot(self):   
        self.figure = plt.figure()
        
        plt.xticks(
            range(0,self.params['steps']*self.params['epochs'], self.params['steps']),
            range(0,self.params['epochs']))
        plt.xlabel('Epoch')
        plt.ylabel('Loss & Metrics ($Log_{10}$ scale)')

        self.plot = {}
        for metric in self.model.metrics_names:
          self.plot[metric], = plt.plot([],[], label = metric)
          
        max_y = [max(log.values()) for log in self.logs_history]
        
        self.title = plt.title(f'batches:0')
        plt.xlim(0,len(self.logs_history)) 
        plt.ylim(0,max(max_y))

           
        plt.legend(loc='upper right')
  
    def animation_function(self,frame):
        batch = frame % self.params['steps']
        self.title.set_text(f'batch:{batch}')
        x = list(range(frame))
        
        for metric in self.model.metrics_names:
            y = [log[metric] for log in self.logs_history[:frame]]
            self.plot[metric].set_data(x,y)
        
    def on_train_batch_end(self, batch, logs=None):
        logarithm_transform = lambda item: (item[0], np.log(item[1]))
        logs = dict(map(logarithm_transform,logs.items()))
        self.logs_history.append(logs)
       
    def on_train_end(self, logs=None):
        self.set_plot()
        num_frames = int(self.duration*self.fps)
        num_batches = self.params['steps']*self.params['epochs']
        selected_batches = range(0, num_batches , num_batches//num_frames )
        interval = 1000*(1/self.fps)
        anim_created = FuncAnimation(self.figure, 
                                     self.animation_function,
                                     frames=selected_batches,
                                     interval=interval)
        video = anim_created.to_html5_video()
        
        html = display.HTML(video)
        display.display(html)
        plt.close()

Usaremos o mesmo modelo de antes, mas com mais exemplos de treinamento:

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

model = keras.Sequential()
model.add(keras.layers.Dense(10, input_dim = 1, activation='relu'))
model.add(keras.layers.Dense(10, activation='relu'))
model.add(keras.layers.Dense(1))
model.compile(
    optimizer=keras.optimizers.RMSprop(learning_rate=0.1),
    loss = "mean_squared_error",
    metrics = ["mean_absolute_error"]
)

def create_sample(sample_size, train_test_proportion = 0.9):
    x = np.random.uniform(low = 0, high = 10, size = sample_size)
    y = x**2
    train_test_split = int(sample_size*train_test_proportion)
    x_train, x_test = (x[:train_test_split],x[train_test_split:])
    y_train, y_test = (y[:train_test_split],y[train_test_split:])
    return (x_train,x_test,y_train,y_test)

x_train,x_test,y_train,y_test = create_sample(35200)


model.fit(
    x_train,
    y_train,
    batch_size=32,
    epochs=2,
    verbose=0,
    callbacks=[TrainingAnimationCallback()],
)

Nossa saída é uma animação das métricas e da função de perda conforme elas mudam durante o processo de treinamento:

Conclusão

Neste guia, analisamos a implementação de retornos de chamada personalizados no Keras.
Existem duas opções para implementar retornos de chamada personalizados – através da subclassificação do keras.callbacks.Callback classe, ou usando o keras.callbacks.LambdaCallback classe.

Vimos um exemplo prático usando LambdaCallbackpara enviar um e-mail no final do loop de treinamento e um exemplo de subclassificação do Callback classe que cria uma animação do loop de treinamento.

Embora Keras tenha muitos retornos de chamada integrados, saber como implementar um retorno de chamada personalizado pode ser útil para aplicativos mais específicos.