将 SessionRunHook 迁移到 Keras 回调

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

在 TensorFlow 1 中,要自定义训练的行为,可以使用 tf.estimator.SessionRunHooktf.estimator.Estimator。本指南演示了如何使用 tf.keras.callbacks.Callback API 从 SessionRunHook 迁移到 TensorFlow 2 的自定义回调,此 API 与 Keras Model.fit 一起用于训练(以及 Model.evaluateModel.predict)。您将通过实现 SessionRunHookCallback 任务来学习如何做到这一点,此任务会在训练期间测量每秒的样本数。

回调的示例为检查点保存 (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
2022-12-14 21:05:28.476156: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:05:28.476264: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:05:28.476274: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
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.trainhooks 参数。

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()])
INFO:tensorflow:Using default config.
WARNING:tensorflow:Using temporary folder as model directory: /tmpfs/tmp/tmpx2ok2ea3
INFO:tensorflow:Using config: {'_model_dir': '/tmpfs/tmp/tmpx2ok2ea3', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/training/training_util.py:396: Variable.initialized_value (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Use Variable.read_value. Variables in 2.X are initialized automatically both in eager and graph (inside tf.defun) contexts.
INFO:tensorflow:Calling model_fn.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/training/adagrad.py:138: calling Constant.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Call initializer instance with the dtype argument instead of passing it to the constructor
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0...
INFO:tensorflow:Saving checkpoints for 0 into /tmpfs/tmp/tmpx2ok2ea3/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0...
Time: 2022-12-14 21:05:33.564674 , Step #: 0 , Examples per second: 2.6388003375708364
INFO:tensorflow:loss = 0.2890536, step = 0
Time: 2022-12-14 21:05:33.597332 , Step #: 10 , Examples per second: 306.1669854154197
Time: 2022-12-14 21:05:33.604084 , Step #: 20 , Examples per second: 1480.9349622201821
Time: 2022-12-14 21:05:33.610736 , Step #: 30 , Examples per second: 1503.173135505143
Time: 2022-12-14 21:05:33.617525 , Step #: 40 , Examples per second: 1473.0294303575192
Time: 2022-12-14 21:05:33.624239 , Step #: 50 , Examples per second: 1489.4545454545455
Time: 2022-12-14 21:05:33.631684 , Step #: 60 , Examples per second: 1343.2088644078653
Time: 2022-12-14 21:05:33.638755 , Step #: 70 , Examples per second: 1414.1281186783547
Time: 2022-12-14 21:05:33.645996 , Step #: 80 , Examples per second: 1381.068159367797
Time: 2022-12-14 21:05:33.653013 , Step #: 90 , Examples per second: 1425.0829029627616
INFO:tensorflow:global_step/sec: 1037.69
Time: 2022-12-14 21:05:33.661947 , Step #: 100 , Examples per second: 1119.3466947772945
INFO:tensorflow:loss = 2.0069905e-05, step = 100 (0.097 sec)
Time: 2022-12-14 21:05:33.669810 , Step #: 110 , Examples per second: 1271.8105461050973
Time: 2022-12-14 21:05:33.676815 , Step #: 120 , Examples per second: 1427.4108358290225
Time: 2022-12-14 21:05:33.683562 , Step #: 130 , Examples per second: 1482.1909675595448
Time: 2022-12-14 21:05:33.690362 , Step #: 140 , Examples per second: 1470.5504522824485
Time: 2022-12-14 21:05:33.697251 , Step #: 150 , Examples per second: 1451.768370772905
Time: 2022-12-14 21:05:33.703805 , Step #: 160 , Examples per second: 1525.5342983923765
Time: 2022-12-14 21:05:33.710547 , Step #: 170 , Examples per second: 1483.2917211868303
Time: 2022-12-14 21:05:33.717187 , Step #: 180 , Examples per second: 1506.141913243321
Time: 2022-12-14 21:05:33.723858 , Step #: 190 , Examples per second: 1498.9293117003788
INFO:tensorflow:global_step/sec: 1418.02
Time: 2022-12-14 21:05:33.732304 , Step #: 200 , Examples per second: 1183.995483415667
INFO:tensorflow:loss = 7.567363e-05, step = 200 (0.070 sec)
Time: 2022-12-14 21:05:33.740206 , Step #: 210 , Examples per second: 1265.555488504013
Time: 2022-12-14 21:05:33.747279 , Step #: 220 , Examples per second: 1413.7467978967238
Time: 2022-12-14 21:05:33.754154 , Step #: 230 , Examples per second: 1454.6382742595547
Time: 2022-12-14 21:05:33.760685 , Step #: 240 , Examples per second: 1531.0472713998904
Time: 2022-12-14 21:05:33.767590 , Step #: 250 , Examples per second: 1448.2093778054002
Time: 2022-12-14 21:05:33.774400 , Step #: 260 , Examples per second: 1468.5424179825636
Time: 2022-12-14 21:05:33.781041 , Step #: 270 , Examples per second: 1505.655311052877
Time: 2022-12-14 21:05:33.787675 , Step #: 280 , Examples per second: 1507.603608784731
Time: 2022-12-14 21:05:33.794272 , Step #: 290 , Examples per second: 1515.776083264067
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 300...
INFO:tensorflow:Saving checkpoints for 300 into /tmpfs/tmp/tmpx2ok2ea3/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 300...
INFO:tensorflow:Loss for final step: 2.7116595e-05.
<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7fae99b2ac40>

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
Time: 2022-12-14 21:05:34.715803 , Step #: 0 , Examples per second: 20.32853849177442
Time: 2022-12-14 21:05:34.735514 , Step #: 10 , Examples per second: 507.25071655762093
Time: 2022-12-14 21:05:34.752205 , Step #: 20 , Examples per second: 599.0921426632958
Time: 2022-12-14 21:05:34.769600 , Step #: 30 , Examples per second: 574.9008319969297
Time: 2022-12-14 21:05:34.786129 , Step #: 40 , Examples per second: 604.9680518094359
Time: 2022-12-14 21:05:34.802761 , Step #: 50 , Examples per second: 601.2477064220184
Time: 2022-12-14 21:05:34.819602 , Step #: 60 , Examples per second: 593.8164880438322
Time: 2022-12-14 21:05:34.836180 , Step #: 70 , Examples per second: 603.210561891476
Time: 2022-12-14 21:05:34.852972 , Step #: 80 , Examples per second: 595.5026763023014
Time: 2022-12-14 21:05:34.869650 , Step #: 90 , Examples per second: 599.5802956228379
Time: 2022-12-14 21:05:34.886414 , Step #: 100 , Examples per second: 596.5359616560709
Time: 2022-12-14 21:05:34.902936 , Step #: 110 , Examples per second: 605.2386724386724
Time: 2022-12-14 21:05:34.919523 , Step #: 120 , Examples per second: 602.889751329596
Time: 2022-12-14 21:05:34.936131 , Step #: 130 , Examples per second: 602.1108240022969
Time: 2022-12-14 21:05:34.952220 , Step #: 140 , Examples per second: 621.5533261214267
Time: 2022-12-14 21:05:34.968538 , Step #: 150 , Examples per second: 612.7991818248229
Time: 2022-12-14 21:05:34.984989 , Step #: 160 , Examples per second: 607.8877648627497
Time: 2022-12-14 21:05:35.001109 , Step #: 170 , Examples per second: 620.3490504644146
Time: 2022-12-14 21:05:35.016755 , Step #: 180 , Examples per second: 639.1222990887758
Time: 2022-12-14 21:05:35.032922 , Step #: 190 , Examples per second: 618.5376788084353
Time: 2022-12-14 21:05:35.049564 , Step #: 200 , Examples per second: 600.8945430581224
Time: 2022-12-14 21:05:35.066073 , Step #: 210 , Examples per second: 605.7281497313846
Time: 2022-12-14 21:05:35.082339 , Step #: 220 , Examples per second: 614.8022631995544
Time: 2022-12-14 21:05:35.098710 , Step #: 230 , Examples per second: 610.8269012320508
Time: 2022-12-14 21:05:35.115086 , Step #: 240 , Examples per second: 610.6490405613953
Time: 2022-12-14 21:05:35.131756 , Step #: 250 , Examples per second: 599.8890128436168
Time: 2022-12-14 21:05:35.148034 , Step #: 260 , Examples per second: 614.325009154156
Time: 2022-12-14 21:05:35.164035 , Step #: 270 , Examples per second: 624.9614828721708
Time: 2022-12-14 21:05:35.180192 , Step #: 280 , Examples per second: 618.9027593330383
Time: 2022-12-14 21:05:35.195872 , Step #: 290 , Examples per second: 637.7811568639377
{'loss': [0.20529791712760925]}

后续步骤

通过下列方式详细了解回调:

此外,您可能还会发现下列与迁移相关的资源十分有用: