将 SessionRunHook 迁移到 Keras 回调

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

在 TensorFlow 1 中，要自定义训练行为，您可以使用 tf.estimator.SessionRunHook 与 tf.estimator.Estimator。本指南演示了如何从 SessionRunHook 迁移到 TensorFlow 2 的自定义回调，使用 tf.keras.callbacks.Callback API，该 API 与 Keras Model.fit 一起用于训练（以及 Model.evaluate 和 Model.predict）。您将学习如何通过实现一个 SessionRunHook 和一个 Callback 任务来完成此操作，该任务在训练期间测量每秒的示例数。

回调示例包括检查点保存 (tf.keras.callbacks.ModelCheckpoint) 和 TensorBoard 摘要写入。Keras 回调 是在内置 Keras Model.fit/Model.evaluate/Model.predict API 中训练/评估/预测的不同阶段调用的对象。您可以在 tf.keras.callbacks.Callback API 文档以及 编写您自己的回调 和 使用内置方法进行训练和评估（“使用回调”部分）指南中了解有关回调的更多信息。

设置

从导入和一个简单的演示数据集开始

import tensorflow as tf
import tensorflow.compat.v1 as tf1

import time
from datetime import datetime
from absl import flags

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

TensorFlow 1：使用 tf.estimator API 创建自定义 SessionRunHook

以下 TensorFlow 1 示例展示了如何设置一个自定义 SessionRunHook，该挂钩在训练期间测量每秒的示例数。在创建挂钩 (LoggerHook) 后，将其传递给 tf.estimator.Estimator.train 的 hooks 参数。

def _input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (features, labels)).batch(1).repeat(100)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

class LoggerHook(tf1.train.SessionRunHook):
  """Logs loss and runtime."""

  def begin(self):
    self._step = -1
    self._start_time = time.time()
    self.log_frequency = 10

  def before_run(self, run_context):
    self._step += 1

  def after_run(self, run_context, run_values):
    if self._step % self.log_frequency == 0:
      current_time = time.time()
      duration = current_time - self._start_time
      self._start_time = current_time
      examples_per_sec = self.log_frequency / duration
      print('Time:', datetime.now(), ', Step #:', self._step,
            ', Examples per second:', examples_per_sec)

estimator = tf1.estimator.Estimator(model_fn=_model_fn)

# Begin training.
estimator.train(_input_fn, hooks=[LoggerHook()])

TensorFlow 2：为 Model.fit 创建自定义 Keras 回调

在 TensorFlow 2 中，当您使用内置 Keras Model.fit（或 Model.evaluate）进行训练/评估时，您可以配置一个自定义 tf.keras.callbacks.Callback，然后将其传递给 Model.fit（或 Model.evaluate）的 callbacks 参数。（在 编写您自己的回调 指南中了解更多信息。）

在下面的示例中，您将编写一个自定义 tf.keras.callbacks.Callback，该回调记录各种指标 - 它将测量每秒的示例数，这应该与之前 SessionRunHook 示例中的指标相当。

class CustomCallback(tf.keras.callbacks.Callback):

    def on_train_begin(self, logs = None):
      self._step = -1
      self._start_time = time.time()
      self.log_frequency = 10

    def on_train_batch_begin(self, batch, logs = None):
      self._step += 1

    def on_train_batch_end(self, batch, logs = None):
      if self._step % self.log_frequency == 0:
        current_time = time.time()
        duration = current_time - self._start_time
        self._start_time = current_time
        examples_per_sec = self.log_frequency / duration
        print('Time:', datetime.now(), ', Step #:', self._step,
              ', Examples per second:', examples_per_sec)

callback = CustomCallback()

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).batch(1).repeat(100)

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer, "mse")

# Begin training.
result = model.fit(dataset, callbacks=[callback], verbose = 0)
# Provide the results of training metrics.
result.history

后续步骤

了解有关回调的更多信息

您可能还会发现以下与迁移相关的资源很有用