將 LoggingTensorHook 和 StopAtStepHook 遷移至 Keras 回呼

在 TensorFlow.org 上查看 在 Google Colab 中執行 在 GitHub 上查看原始碼 下載筆記本

在 TensorFlow 1 中，您可以使用 tf.estimator.LoggingTensorHook 來監控和記錄張量，而使用 tf.estimator.Estimator 進行訓練時，tf.estimator.StopAtStepHook 可協助在指定步驟停止訓練。本筆記本示範如何使用自訂 Keras 回呼 (tf.keras.callbacks.Callback) 和 Model.fit，將這些 API 從 TensorFlow 1 遷移至 TensorFlow 2 中的對等項目。

Keras 回呼是在內建 Keras Model.fit/Model.evaluate/Model.predict API 中，於訓練/評估/預測期間不同時間點呼叫的物件。您可以在 tf.keras.callbacks.Callback API 文件，以及撰寫您自己的回呼和使用內建方法進行訓練和評估 (「使用回呼」章節) 指南中進一步瞭解回呼。如要將 TensorFlow 1 中的 SessionRunHook 遷移至 TensorFlow 2 中的 Keras 回呼，請參閱使用輔助邏輯遷移訓練作業指南。

設定

從匯入項目和簡單資料集開始，以進行示範

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

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]

# Define an input function.
def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

TensorFlow 1：使用 tf.estimator API 記錄張量並停止訓練

在 TensorFlow 1 中，您可以定義各種掛鉤來控制訓練行為。然後，將這些掛鉤傳遞至 tf.estimator.EstimatorSpec。

在以下範例中

  • 如要監控/記錄張量 (例如模型權重或損失)，您可以使用 tf.estimator.LoggingTensorHook (tf.train.LoggingTensorHook 是其別名)。
  • 如要在特定步驟停止訓練，您可以使用 tf.estimator.StopAtStepHook (tf.train.StopAtStepHook 是其別名)。
def _model_fn(features, labels, mode):
  dense = tf1.layers.Dense(1)
  logits = dense(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())

  # Define the stop hook.
  stop_hook = tf1.train.StopAtStepHook(num_steps=2)

  # Access tensors to be logged by names.
  kernel_name = tf.identity(dense.weights[0])
  bias_name = tf.identity(dense.weights[1])
  logging_weight_hook = tf1.train.LoggingTensorHook(
      tensors=[kernel_name, bias_name],
      every_n_iter=1)
  # Log the training loss by the tensor object.
  logging_loss_hook = tf1.train.LoggingTensorHook(
      {'loss from LoggingTensorHook': loss},
      every_n_secs=3)

  # Pass all hooks to `EstimatorSpec`.
  return tf1.estimator.EstimatorSpec(mode,
                                     loss=loss,
                                     train_op=train_op,
                                     training_hooks=[stop_hook,
                                                     logging_weight_hook,
                                                     logging_loss_hook])

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

# Begin training.
# The training will stop after 2 steps, and the weights/loss will also be logged.
estimator.train(_input_fn)

TensorFlow 2：使用自訂回呼和 Model.fit 記錄張量並停止訓練

在 TensorFlow 2 中，當您使用內建 Keras Model.fit (或 Model.evaluate) 進行訓練/評估時，可以透過定義自訂 Keras tf.keras.callbacks.Callback 來設定張量監控和停止訓練。然後，將這些回呼傳遞至 Model.fit (或 Model.evaluate) 的 callbacks 參數。(詳情請參閱撰寫您自己的回呼指南。)

在以下範例中

  • 如要重新建立 StopAtStepHook 的功能，請定義自訂回呼 (以下命名為 StopAtStepCallback)，在其中覆寫 on_batch_end 方法，以便在特定步驟數後停止訓練。
  • 如要重新建立 LoggingTensorHook 行為，請定義自訂回呼 (LoggingTensorCallback)，在其中手動記錄和輸出記錄的張量，因為不支援依名稱存取張量。您也可以在自訂回呼內實作記錄頻率。以下範例會每兩個步驟列印權重。其他策略 (例如每 N 秒記錄一次) 也可行。
class StopAtStepCallback(tf.keras.callbacks.Callback):
  def __init__(self, stop_step=None):
    super().__init__()
    self._stop_step = stop_step

  def on_batch_end(self, batch, logs=None):
    if self.model.optimizer.iterations >= self._stop_step:
      self.model.stop_training = True
      print('\nstop training now')

class LoggingTensorCallback(tf.keras.callbacks.Callback):
  def __init__(self, every_n_iter):
      super().__init__()
      self._every_n_iter = every_n_iter
      self._log_count = every_n_iter

  def on_batch_end(self, batch, logs=None):
    if self._log_count > 0:
      self._log_count -= 1
      print("Logging Tensor Callback: dense/kernel:",
            model.layers[0].weights[0])
      print("Logging Tensor Callback: dense/bias:",
            model.layers[0].weights[1])
      print("Logging Tensor Callback loss:", logs["loss"])
    else:
      self._log_count -= self._every_n_iter

完成後，將新的回呼 (亦即 StopAtStepCallback 和 LoggingTensorCallback) 傳遞至 Model.fit 的 callbacks 參數。

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
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.
# The training will stop after 2 steps, and the weights/loss will also be logged.
model.fit(dataset, callbacks=[StopAtStepCallback(stop_step=2),
                              LoggingTensorCallback(every_n_iter=2)])

後續步驟

進一步瞭解回呼，請參閱：

您可能也會覺得下列與遷移相關的資源很有用：