编写自己的回调函数

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

简介

回调函数是一种可以在训练、评估或推断过程中自定义 Keras 模型行为的强大工具。示例包括使用 TensorBoard 来可视化训练进度和结果的 tf.keras.callbacks.TensorBoard,以及用来在训练期间定期保存模型的 tf.keras.callbacks.ModelCheckpoint

在本指南中,您将了解什么是 Keras 回调函数,它可以做什么,以及如何构建自己的回调函数。我们提供了一些简单回调函数应用的演示,以帮助您入门。

设置

import tensorflow as tf
from tensorflow import keras

Keras 回调函数概述

所有回调函数都将 keras.callbacks.Callback 类作为子类,并重写在训练、测试和预测的各个阶段调用的一组方法。回调函数对于在训练期间了解模型的内部状态和统计信息十分有用。

您可以将回调函数的列表(作为关键字参数 callbacks)传递给以下模型方法:

回调函数方法概述

全局方法

on_(train|test|predict)_begin(self, logs=None)

fit/evaluate/predict 开始时调用。

on_(train|test|predict)_end(self, logs=None)

fit/evaluate/predict 结束时调用。

Batch-level methods for training/testing/predicting

on_(train|test|predict)_batch_begin(self, batch, logs=None)

正好在训练/测试/预测期间处理批次之前调用。

on_(train|test|predict)_batch_end(self, batch, logs=None)

在训练/测试/预测批次结束时调用。在此方法中,logs 是包含指标结果的字典。

周期级方法(仅训练)

on_epoch_begin(self, epoch, logs=None)

在训练期间周期开始时调用。

on_epoch_end(self, epoch, logs=None)

在训练期间周期开始时调用。

基本示例

让我们来看一个具体的例子。首先,导入 Tensorflow 并定义一个简单的序列式 Keras 模型:

# Define the Keras model to add callbacks to
def get_model():
    model = keras.Sequential()
    model.add(keras.layers.Dense(1, input_dim=784))
    model.compile(
        optimizer=keras.optimizers.RMSprop(learning_rate=0.1),
        loss="mean_squared_error",
        metrics=["mean_absolute_error"],
    )
    return model

然后,从 Keras 数据集 API 加载 MNIST 数据进行训练和测试:

# Load example MNIST data and pre-process it
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
x_test = x_test.reshape(-1, 784).astype("float32") / 255.0

# Limit the data to 1000 samples
x_train = x_train[:1000]
y_train = y_train[:1000]
x_test = x_test[:1000]
y_test = y_test[:1000]
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step

接下来,定义一个简单的自定义回调函数来记录以下内容:

  • fit/evaluate/predict 开始和结束的时间
  • 每个周期开始和结束的时间
  • 每个训练批次开始和结束的时间
  • 每个评估(测试)批次开始和结束的时间
  • 每次推断(预测)批次开始和结束的时间
class CustomCallback(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        keys = list(logs.keys())
        print("Starting training; got log keys: {}".format(keys))

    def on_train_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop training; got log keys: {}".format(keys))

    def on_epoch_begin(self, epoch, logs=None):
        keys = list(logs.keys())
        print("Start epoch {} of training; got log keys: {}".format(epoch, keys))

    def on_epoch_end(self, epoch, logs=None):
        keys = list(logs.keys())
        print("End epoch {} of training; got log keys: {}".format(epoch, keys))

    def on_test_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start testing; got log keys: {}".format(keys))

    def on_test_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop testing; got log keys: {}".format(keys))

    def on_predict_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start predicting; got log keys: {}".format(keys))

    def on_predict_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop predicting; got log keys: {}".format(keys))

    def on_train_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Training: start of batch {}; got log keys: {}".format(batch, keys))

    def on_train_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Training: end of batch {}; got log keys: {}".format(batch, keys))

    def on_test_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Evaluating: start of batch {}; got log keys: {}".format(batch, keys))

    def on_test_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Evaluating: end of batch {}; got log keys: {}".format(batch, keys))

    def on_predict_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Predicting: start of batch {}; got log keys: {}".format(batch, keys))

    def on_predict_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print("...Predicting: end of batch {}; got log keys: {}".format(batch, keys))

我们来试一下:

model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=128,
    epochs=1,
    verbose=0,
    validation_split=0.5,
    callbacks=[CustomCallback()],
)

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

res = model.predict(x_test, batch_size=128, callbacks=[CustomCallback()])
Starting training; got log keys: []
Start epoch 0 of training; got log keys: []
...Training: start of batch 0; got log keys: []
...Training: end of batch 0; got log keys: ['loss', 'mean_absolute_error']
...Training: start of batch 1; got log keys: []
...Training: end of batch 1; got log keys: ['loss', 'mean_absolute_error']
...Training: start of batch 2; got log keys: []
...Training: end of batch 2; got log keys: ['loss', 'mean_absolute_error']
...Training: start of batch 3; got log keys: []
...Training: end of batch 3; got log keys: ['loss', 'mean_absolute_error']
Start testing; got log keys: []
...Evaluating: start of batch 0; got log keys: []
...Evaluating: end of batch 0; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 1; got log keys: []
...Evaluating: end of batch 1; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 2; got log keys: []
...Evaluating: end of batch 2; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 3; got log keys: []
...Evaluating: end of batch 3; got log keys: ['loss', 'mean_absolute_error']
Stop testing; got log keys: ['loss', 'mean_absolute_error']
End epoch 0 of training; got log keys: ['loss', 'mean_absolute_error', 'val_loss', 'val_mean_absolute_error']
Stop training; got log keys: ['loss', 'mean_absolute_error', 'val_loss', 'val_mean_absolute_error']
Start testing; got log keys: []
...Evaluating: start of batch 0; got log keys: []
...Evaluating: end of batch 0; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 1; got log keys: []
...Evaluating: end of batch 1; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 2; got log keys: []
...Evaluating: end of batch 2; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 3; got log keys: []
...Evaluating: end of batch 3; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 4; got log keys: []
...Evaluating: end of batch 4; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 5; got log keys: []
...Evaluating: end of batch 5; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 6; got log keys: []
...Evaluating: end of batch 6; got log keys: ['loss', 'mean_absolute_error']
...Evaluating: start of batch 7; got log keys: []
...Evaluating: end of batch 7; got log keys: ['loss', 'mean_absolute_error']
Stop testing; got log keys: ['loss', 'mean_absolute_error']
Start predicting; got log keys: []
...Predicting: start of batch 0; got log keys: []
...Predicting: end of batch 0; got log keys: ['outputs']
...Predicting: start of batch 1; got log keys: []
...Predicting: end of batch 1; got log keys: ['outputs']
...Predicting: start of batch 2; got log keys: []
...Predicting: end of batch 2; got log keys: ['outputs']
...Predicting: start of batch 3; got log keys: []
...Predicting: end of batch 3; got log keys: ['outputs']
...Predicting: start of batch 4; got log keys: []
...Predicting: end of batch 4; got log keys: ['outputs']
...Predicting: start of batch 5; got log keys: []
...Predicting: end of batch 5; got log keys: ['outputs']
...Predicting: start of batch 6; got log keys: []
...Predicting: end of batch 6; got log keys: ['outputs']
...Predicting: start of batch 7; got log keys: []
...Predicting: end of batch 7; got log keys: ['outputs']
Stop predicting; got log keys: []

logs 字典的用法

logs 字典包含损失值,以及批次或周期结束时的所有指标。示例包括损失和平均绝对误差。

class LossAndErrorPrintingCallback(keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        print("For batch {}, loss is {:7.2f}.".format(batch, logs["loss"]))

    def on_test_batch_end(self, batch, logs=None):
        print("For batch {}, loss is {:7.2f}.".format(batch, logs["loss"]))

    def on_epoch_end(self, epoch, logs=None):
        print(
            "The average loss for epoch {} is {:7.2f} "
            "and mean absolute error is {:7.2f}.".format(
                epoch, logs["loss"], logs["mean_absolute_error"]
            )
        )


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=128,
    epochs=2,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback()],
)

res = model.evaluate(
    x_test,
    y_test,
    batch_size=128,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback()],
)
For batch 0, loss is   24.99.
For batch 1, loss is  494.81.
For batch 2, loss is  338.22.
For batch 3, loss is  256.11.
For batch 4, loss is  206.34.
For batch 5, loss is  172.97.
For batch 6, loss is  149.13.
For batch 7, loss is  134.18.
The average loss for epoch 0 is  134.18 and mean absolute error is    6.11.
For batch 0, loss is    4.78.
For batch 1, loss is    5.10.
For batch 2, loss is    4.68.
For batch 3, loss is    4.67.
For batch 4, loss is    4.59.
For batch 5, loss is    4.66.
For batch 6, loss is    4.65.
For batch 7, loss is    4.60.
The average loss for epoch 1 is    4.60 and mean absolute error is    1.75.
For batch 0, loss is    5.27.
For batch 1, loss is    4.53.
For batch 2, loss is    4.58.
For batch 3, loss is    4.52.
For batch 4, loss is    4.67.
For batch 5, loss is    4.66.
For batch 6, loss is    4.60.
For batch 7, loss is    4.54.

self.model 属性的用法

除了在调用其中一种方法时接收日志信息外,回调还可以访问与当前一轮训练/评估/推断有关的模型:self.model

以下是您可以在回调函数中使用 self.model 进行的一些操作:

  • 设置 self.model.stop_training = True 以立即中断训练。
  • 转变优化器(可作为 self.model.optimizer)的超参数,例如 self.model.optimizer.learning_rate
  • 定期保存模型。
  • 在每个周期结束时,在少量测试样本上记录 model.predict() 的输出,以用作训练期间的健全性检查。
  • 在每个周期结束时提取中间特征的可视化,随时间推移监视模型当前的学习内容。
  • 其他

下面我们通过几个示例来看看它是如何工作的。

Keras 回调函数应用示例

在达到最小损失时尽早停止

第一个示例展示了如何通过设置 self.model.stop_training(布尔)属性来创建能够在达到最小损失时停止训练的 Callback。您还可以提供参数 patience 来指定在达到局部最小值后应该等待多少个周期然后停止。

tf.keras.callbacks.EarlyStopping 提供了一种更完整、更通用的实现。

import numpy as np


class EarlyStoppingAtMinLoss(keras.callbacks.Callback):
    """Stop training when the loss is at its min, i.e. the loss stops decreasing.

  Arguments:
      patience: Number of epochs to wait after min has been hit. After this
      number of no improvement, training stops.
  """

    def __init__(self, patience=0):
        super(EarlyStoppingAtMinLoss, self).__init__()
        self.patience = patience
        # best_weights to store the weights at which the minimum loss occurs.
        self.best_weights = None

    def on_train_begin(self, logs=None):
        # The number of epoch it has waited when loss is no longer minimum.
        self.wait = 0
        # The epoch the training stops at.
        self.stopped_epoch = 0
        # Initialize the best as infinity.
        self.best = np.Inf

    def on_epoch_end(self, epoch, logs=None):
        current = logs.get("loss")
        if np.less(current, self.best):
            self.best = current
            self.wait = 0
            # Record the best weights if current results is better (less).
            self.best_weights = self.model.get_weights()
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.model.stop_training = True
                print("Restoring model weights from the end of the best epoch.")
                self.model.set_weights(self.best_weights)

    def on_train_end(self, logs=None):
        if self.stopped_epoch > 0:
            print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=64,
    steps_per_epoch=5,
    epochs=30,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback(), EarlyStoppingAtMinLoss()],
)
For batch 0, loss is   28.00.
For batch 1, loss is  421.44.
For batch 2, loss is  287.19.
For batch 3, loss is  218.01.
For batch 4, loss is  175.96.
The average loss for epoch 0 is  175.96 and mean absolute error is    7.89.
For batch 0, loss is    5.78.
For batch 1, loss is    6.13.
For batch 2, loss is    6.36.
For batch 3, loss is    5.85.
For batch 4, loss is    5.86.
The average loss for epoch 1 is    5.86 and mean absolute error is    2.00.
For batch 0, loss is    4.92.
For batch 1, loss is    4.27.
For batch 2, loss is    4.53.
For batch 3, loss is    4.59.
For batch 4, loss is    4.43.
The average loss for epoch 2 is    4.43 and mean absolute error is    1.69.
For batch 0, loss is    5.66.
For batch 1, loss is    5.16.
For batch 2, loss is    4.60.
For batch 3, loss is    4.96.
For batch 4, loss is    6.47.
The average loss for epoch 3 is    6.47 and mean absolute error is    2.02.
Restoring model weights from the end of the best epoch.
Epoch 00004: early stopping
<tensorflow.python.keras.callbacks.History at 0x7f1968109da0>

学习率规划

在此示例中,我们展示了如何在学习过程中使用自定义回调来动态更改优化器的学习率。

有关更通用的实现,请查看 callbacks.LearningRateScheduler

class CustomLearningRateScheduler(keras.callbacks.Callback):
    """Learning rate scheduler which sets the learning rate according to schedule.

  Arguments:
      schedule: a function that takes an epoch index
          (integer, indexed from 0) and current learning rate
          as inputs and returns a new learning rate as output (float).
  """

    def __init__(self, schedule):
        super(CustomLearningRateScheduler, self).__init__()
        self.schedule = schedule

    def on_epoch_begin(self, epoch, logs=None):
        if not hasattr(self.model.optimizer, "lr"):
            raise ValueError('Optimizer must have a "lr" attribute.')
        # Get the current learning rate from model's optimizer.
        lr = float(tf.keras.backend.get_value(self.model.optimizer.learning_rate))
        # Call schedule function to get the scheduled learning rate.
        scheduled_lr = self.schedule(epoch, lr)
        # Set the value back to the optimizer before this epoch starts
        tf.keras.backend.set_value(self.model.optimizer.lr, scheduled_lr)
        print("\nEpoch %05d: Learning rate is %6.4f." % (epoch, scheduled_lr))


LR_SCHEDULE = [
    # (epoch to start, learning rate) tuples
    (3, 0.05),
    (6, 0.01),
    (9, 0.005),
    (12, 0.001),
]


def lr_schedule(epoch, lr):
    """Helper function to retrieve the scheduled learning rate based on epoch."""
    if epoch < LR_SCHEDULE[0][0] or epoch > LR_SCHEDULE[-1][0]:
        return lr
    for i in range(len(LR_SCHEDULE)):
        if epoch == LR_SCHEDULE[i][0]:
            return LR_SCHEDULE[i][1]
    return lr


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=64,
    steps_per_epoch=5,
    epochs=15,
    verbose=0,
    callbacks=[
        LossAndErrorPrintingCallback(),
        CustomLearningRateScheduler(lr_schedule),
    ],
)
Epoch 00000: Learning rate is 0.1000.
For batch 0, loss is   26.58.
For batch 1, loss is  494.68.
For batch 2, loss is  339.22.
For batch 3, loss is  257.81.
For batch 4, loss is  207.42.
The average loss for epoch 0 is  207.42 and mean absolute error is    8.60.

Epoch 00001: Learning rate is 0.1000.
For batch 0, loss is    6.30.
For batch 1, loss is    6.45.
For batch 2, loss is    6.02.
For batch 3, loss is    5.85.
For batch 4, loss is    5.73.
The average loss for epoch 1 is    5.73 and mean absolute error is    1.97.

Epoch 00002: Learning rate is 0.1000.
For batch 0, loss is    5.97.
For batch 1, loss is    6.09.
For batch 2, loss is    5.57.
For batch 3, loss is    5.10.
For batch 4, loss is    5.02.
The average loss for epoch 2 is    5.02 and mean absolute error is    1.82.

Epoch 00003: Learning rate is 0.0500.
For batch 0, loss is    3.41.
For batch 1, loss is    3.07.
For batch 2, loss is    3.92.
For batch 3, loss is    3.89.
For batch 4, loss is    3.88.
The average loss for epoch 3 is    3.88 and mean absolute error is    1.60.

Epoch 00004: Learning rate is 0.0500.
For batch 0, loss is    4.15.
For batch 1, loss is    4.11.
For batch 2, loss is    4.39.
For batch 3, loss is    4.00.
For batch 4, loss is    3.76.
The average loss for epoch 4 is    3.76 and mean absolute error is    1.53.

Epoch 00005: Learning rate is 0.0500.
For batch 0, loss is    3.04.
For batch 1, loss is    3.79.
For batch 2, loss is    4.17.
For batch 3, loss is    4.15.
For batch 4, loss is    4.28.
The average loss for epoch 5 is    4.28 and mean absolute error is    1.68.

Epoch 00006: Learning rate is 0.0100.
For batch 0, loss is    3.38.
For batch 1, loss is    3.37.
For batch 2, loss is    3.31.
For batch 3, loss is    3.30.
For batch 4, loss is    3.10.
The average loss for epoch 6 is    3.10 and mean absolute error is    1.39.

Epoch 00007: Learning rate is 0.0100.
For batch 0, loss is    2.34.
For batch 1, loss is    3.27.
For batch 2, loss is    3.12.
For batch 3, loss is    3.20.
For batch 4, loss is    3.36.
The average loss for epoch 7 is    3.36 and mean absolute error is    1.43.

Epoch 00008: Learning rate is 0.0100.
For batch 0, loss is    3.17.
For batch 1, loss is    3.36.
For batch 2, loss is    3.25.
For batch 3, loss is    3.46.
For batch 4, loss is    3.30.
The average loss for epoch 8 is    3.30 and mean absolute error is    1.42.

Epoch 00009: Learning rate is 0.0050.
For batch 0, loss is    4.70.
For batch 1, loss is    4.53.
For batch 2, loss is    4.17.
For batch 3, loss is    4.09.
For batch 4, loss is    4.21.
The average loss for epoch 9 is    4.21 and mean absolute error is    1.57.

Epoch 00010: Learning rate is 0.0050.
For batch 0, loss is    3.18.
For batch 1, loss is    2.82.
For batch 2, loss is    2.83.
For batch 3, loss is    2.96.
For batch 4, loss is    3.04.
The average loss for epoch 10 is    3.04 and mean absolute error is    1.35.

Epoch 00011: Learning rate is 0.0050.
For batch 0, loss is    2.83.
For batch 1, loss is    2.75.
For batch 2, loss is    3.21.
For batch 3, loss is    3.14.
For batch 4, loss is    3.11.
The average loss for epoch 11 is    3.11 and mean absolute error is    1.35.

Epoch 00012: Learning rate is 0.0010.
For batch 0, loss is    4.10.
For batch 1, loss is    3.88.
For batch 2, loss is    3.50.
For batch 3, loss is    3.34.
For batch 4, loss is    3.14.
The average loss for epoch 12 is    3.14 and mean absolute error is    1.39.

Epoch 00013: Learning rate is 0.0010.
For batch 0, loss is    2.93.
For batch 1, loss is    3.17.
For batch 2, loss is    2.94.
For batch 3, loss is    2.86.
For batch 4, loss is    2.90.
The average loss for epoch 13 is    2.90 and mean absolute error is    1.33.

Epoch 00014: Learning rate is 0.0010.
For batch 0, loss is    2.94.
For batch 1, loss is    3.55.
For batch 2, loss is    3.35.
For batch 3, loss is    3.29.
For batch 4, loss is    3.44.
The average loss for epoch 14 is    3.44 and mean absolute error is    1.38.
<tensorflow.python.keras.callbacks.History at 0x7f1968055dd8>

内置 Keras 回调函数

请务必阅读 API 文档查看现有的 Keras 回调函数。应用包括记录到 CSV、保存模型、在 TensorBoard 中可视化指标等等!