<jupyter_start><jupyter_text>Copyright 2021 The TensorFlow Authors.<jupyter_code>#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.<jupyter_output><empty_output><jupyter_text>Migrate SessionRunHook to Keras callbacks              View on TensorFlow.org                Run in Google Colab                View source on GitHub        Download notebook   In TensorFlow 1, to customize the behavior of training, you use `tf.estimator.SessionRunHook` with `tf.estimator.Estimator`. This guide demonstrates how to migrate from `SessionRunHook` to TensorFlow 2's custom callbacks with the `tf.keras.callbacks.Callback` API, which works with Keras `Model.fit` for training (as well as `Model.evaluate` and `Model.predict`). You will learn how to do this by implementing a `SessionRunHook` and a `Callback` task that measures examples per second during training.Examples of callbacks are checkpoint saving (`tf.keras.callbacks.ModelCheckpoint`) and [TensorBoard](`tf.keras.callbacks.TensorBoard`) summary writing. Keras [callbacks](https://www.tensorflow.org/guide/keras/custom_callback) are objects that are called at different points during training/evaluation/prediction in the built-in Keras `Model.fit`/`Model.evaluate`/`Model.predict` APIs. You can learn more about callbacks in the `tf.keras.callbacks.Callback` API docs, as well as the [Writing your own callbacks](https://www.tensorflow.org/guide/keras/custom_callback.ipynb/) and [Training and evaluation with the built-in methods](https://www.tensorflow.org/guide/keras/train_and_evaluate) (the *Using callbacks* section) guides.  SetupStart with imports and a simple dataset for demonstration purposes:<jupyter_code>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.]]<jupyter_output><empty_output><jupyter_text>TensorFlow 1: Create a custom SessionRunHook with tf.estimator APIsThe following TensorFlow 1 examples show how to set up a custom `SessionRunHook` that measures examples per second during training. After creating the hook (`LoggerHook`), pass it to the `hooks` parameter of `tf.estimator.Estimator.train`.<jupyter_code>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()])<jupyter_output><empty_output><jupyter_text>TensorFlow 2: Create a custom Keras callback for Model.fitIn TensorFlow 2, when you use the built-in Keras `Model.fit` (or `Model.evaluate`) for training/evaluation, you can configure a custom `tf.keras.callbacks.Callback`, which you then pass to the `callbacks` parameter of `Model.fit` (or `Model.evaluate`). (Learn more in the [Writing your own callbacks](../..guide/keras/custom_callback.ipynb) guide.)In the example below, you will write a custom `tf.keras.callbacks.Callback` that logs various metrics—it will measure examples per second, which should be comparable to the metrics in the previous `SessionRunHook` example.<jupyter_code>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<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10809883321894304, 'fraction_numerical': 0.013726835964310227, 'mean_word_length': 4.525118483412323, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 4, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '641180', 'license': 'Apache-2.0', 'n_tokens_mistral': 1819, 'n_tokens_neox': 1747, 'n_words': 538, 'path': 'site/en-snapshot/guide/migrate/sessionrunhook_callback.ipynb', 'repo_name': 'phoenix-fork-tensorflow/docs-l10n', 'size': '5828'}