BEGIN_PUBLIC

Resubmission of Keras Tensorboard Callback - enable metrics logging in Eager
END_PUBLIC

Previously a test added in this CL was causing Windows cmake build to fail, I
fixed the test.

*** Reason for rollback ***

Fixed test execution

*** Original change description ***

Automated rollback of commit d4cb01f242dc3ff0f7b0aae7284def46281755f2

PiperOrigin-RevId: 206606179

2 files changed, 105 insertions, 28 deletions

diff --git a/tensorflow/python/keras/callbacks.py b/tensorflow/python/keras/callbacks.py
index d1b9dc27bd..070d41147d 100644
--- a/tensorflow/python/keras/callbacks.py
+++ b/tensorflow/python/keras/callbacks.py
@@ -31,12 +31,14 @@ import time
 import numpy as np
 import six
 
+from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras.engine.training_utils import standardize_input_data
 from tensorflow.python.keras.utils.generic_utils import Progbar
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import summary_ops_v2
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.summary import summary as tf_summary
@@ -716,6 +718,15 @@ class TensorBoard(Callback):
           `embeddings_layer_names`. Numpy array (if the model has a single
           input) or list of Numpy arrays (if the model has multiple inputs).
           Learn [more about embeddings](https://www.tensorflow.org/programmers_guide/embedding)
+
+  Raises:
+      ValueError: If histogram_freq is set and no validation data is provided.
+
+  @compatbility(eager)
+  Using `Tensorboard` callback will work while eager execution is enabled,
+  however outputting histogram summaries of weights and gradients is not
+  supported, and thus `histogram_freq` will be ignored.
+  @end_compatibility
   """
 
   # pylint: enable=line-too-long
@@ -734,6 +745,11 @@ class TensorBoard(Callback):
     super(TensorBoard, self).__init__()
     self.log_dir = log_dir
     self.histogram_freq = histogram_freq
+    if self.histogram_freq and context.executing_eagerly():
+      logging.warning(
+          UserWarning('Weight and gradient histograms not supported for eager'
+                      'execution, setting `histogram_freq` to `0`.'))
+      self.histogram_freq = 0
     self.merged = None
     self.write_graph = write_graph
     self.write_grads = write_grads
@@ -741,18 +757,22 @@ class TensorBoard(Callback):
     self.batch_size = batch_size
     self._current_batch = 0
     self._total_batches_seen = 0
-    # abstracted writer class to be able to stub for testing
-    self._writer_class = tf_summary.FileWriter
     self.embeddings_freq = embeddings_freq
     self.embeddings_layer_names = embeddings_layer_names
     self.embeddings_metadata = embeddings_metadata
     self.embeddings_data = embeddings_data
 
-  def set_model(self, model):
-    """Sets Keras model and creates summary ops."""
+  def _init_writer(self):
+    """Sets file writer."""
+    if context.executing_eagerly():
+      self.writer = summary_ops_v2.create_file_writer(self.log_dir)
+    elif self.write_graph:
+      self.writer = tf_summary.FileWriter(self.log_dir, K.get_session().graph)
+    else:
+      self.writer = tf_summary.FileWriter(self.log_dir)
 
-    self.model = model
-    self.sess = K.get_session()
+  def _make_histogram_ops(self, model):
+    """Defines histogram ops when histogram_freq > 0."""
     # only make histogram summary op if it hasn't already been made
     if self.histogram_freq and self.merged is None:
       for layer in self.model.layers:
@@ -793,8 +813,10 @@ class TensorBoard(Callback):
             def is_indexed_slices(grad):
               return type(grad).__name__ == 'IndexedSlices'
 
-            grads = [grad.values if is_indexed_slices(grad) else grad
-                     for grad in grads]
+            grads = [
+                grad.values if is_indexed_slices(grad) else grad
+                for grad in grads
+            ]
             tf_summary.histogram('{}_grad'.format(mapped_weight_name), grads)
 
         if hasattr(layer, 'output'):
@@ -803,12 +825,16 @@ class TensorBoard(Callback):
               tf_summary.histogram('{}_out_{}'.format(layer.name, i), output)
           else:
             tf_summary.histogram('{}_out'.format(layer.name), layer.output)
-    self.merged = tf_summary.merge_all()
 
-    if self.write_graph:
-      self.writer = self._writer_class(self.log_dir, self.sess.graph)
-    else:
-      self.writer = self._writer_class(self.log_dir)
+  def set_model(self, model):
+    """Sets Keras model and creates summary ops."""
+
+    self.model = model
+    self._init_writer()
+    # histogram summaries only enabled in graph mode
+    if not context.executing_eagerly():
+      self._make_histogram_ops(model)
+      self.merged = tf_summary.merge_all()
 
     # If both embedding_freq and embeddings_data are available, we will
     # visualize embeddings.
@@ -894,17 +920,24 @@ class TensorBoard(Callback):
 
     """
     logs = logs or {}
-    for name, value in logs.items():
-      summary = tf_summary.Summary()
-      summary_value = summary.value.add()
-      summary_value.simple_value = value.item()
-      summary_value.tag = name
-      self.writer.add_summary(summary, step)
+    if context.executing_eagerly():
+      # use v2 summary ops
+      with self.writer.as_default(), summary_ops_v2.always_record_summaries():
+        for name, value in logs.items():
+          summary_ops_v2.scalar(name, value.item(), step=step)
+    else:
+      # use FileWriter from v1 summary
+      for name, value in logs.items():
+        summary = tf_summary.Summary()
+        summary_value = summary.value.add()
+        summary_value.simple_value = value.item()
+        summary_value.tag = name
+        self.writer.add_summary(summary, step)
     self.writer.flush()
 
   def on_train_begin(self, logs=None):
     """Checks if histogram summaries can be run."""
-
+    # will never be set when in eager
     if self.histogram_freq:
       if 'validation_steps' in self.params:
         self._validation_batches = self.params['validation_steps']
diff --git a/tensorflow/python/keras/callbacks_test.py b/tensorflow/python/keras/callbacks_test.py
index 7d830078ce..ffec197028 100644
--- a/tensorflow/python/keras/callbacks_test.py
+++ b/tensorflow/python/keras/callbacks_test.py
@@ -22,6 +22,7 @@ import csv
 import os
 import re
 import shutil
+import tempfile
 import threading
 import unittest
 
@@ -29,10 +30,12 @@ import numpy as np
 
 from tensorflow.core.framework import summary_pb2
 from tensorflow.python import keras
+from tensorflow.python.framework import test_util
 from tensorflow.python.keras import testing_utils
 from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.summary.writer import writer_cache
+from tensorflow.python.training import adam
 
 try:
   import h5py  # pylint:disable=g-import-not-at-top
@@ -917,6 +920,9 @@ class KerasCallbacksTest(test.TestCase):
       def close(self):
         pass
 
+    def _init_writer(obj):
+      obj.writer = FileWriterStub(obj.log_dir)
+
     np.random.seed(1337)
     tmpdir = self.get_temp_dir()
     self.addCleanup(shutil.rmtree, tmpdir)
@@ -940,13 +946,13 @@ class KerasCallbacksTest(test.TestCase):
           loss='categorical_crossentropy',
           optimizer='sgd',
           metrics=['accuracy'])
+      keras.callbacks.TensorBoard._init_writer = _init_writer
       tsb = keras.callbacks.TensorBoard(
           log_dir=tmpdir,
           histogram_freq=1,
           write_images=True,
           write_grads=True,
           batch_size=5)
-      tsb._writer_class = FileWriterStub
       cbks = [tsb]
 
       # fit with validation data
@@ -1118,11 +1124,11 @@ class KerasCallbacksTest(test.TestCase):
       def close(self):
         pass
 
-    logdir = 'fake_dir'
+    temp_dir = self.get_temp_dir()
+    self.addCleanup(shutil.rmtree, temp_dir)
 
-    # log every batch
-    tb_cbk = keras.callbacks.TensorBoard(logdir)
-    tb_cbk.writer = FileWriterStub(logdir)
+    tb_cbk = keras.callbacks.TensorBoard(temp_dir)
+    tb_cbk.writer = FileWriterStub(temp_dir)
 
     for batch in range(5):
       tb_cbk.on_batch_end(batch, {'acc': np.float32(batch)})
@@ -1150,10 +1156,11 @@ class KerasCallbacksTest(test.TestCase):
       def close(self):
         pass
 
-    logdir = 'fake_dir'
+    temp_dir = self.get_temp_dir()
+    self.addCleanup(shutil.rmtree, temp_dir)
 
-    tb_cbk = keras.callbacks.TensorBoard(logdir)
-    tb_cbk.writer = FileWriterStub(logdir)
+    tb_cbk = keras.callbacks.TensorBoard(temp_dir)
+    tb_cbk.writer = FileWriterStub(temp_dir)
 
     tb_cbk.on_batch_end(0, {'acc': np.float32(5.0)})
     tb_cbk.on_epoch_end(0, {'acc': np.float32(10.0)})
@@ -1164,6 +1171,43 @@ class KerasCallbacksTest(test.TestCase):
     self.assertEqual(epoch_step, 0)
     self.assertEqual(epoch_summary.value[0].simple_value, 10.0)
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_Tensorboard_eager(self):
+    with self.test_session():
+      temp_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
+      self.addCleanup(shutil.rmtree, temp_dir)
+
+      (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data(
+          train_samples=TRAIN_SAMPLES,
+          test_samples=TEST_SAMPLES,
+          input_shape=(INPUT_DIM,),
+          num_classes=NUM_CLASSES)
+      y_test = keras.utils.to_categorical(y_test)
+      y_train = keras.utils.to_categorical(y_train)
+
+      model = keras.models.Sequential()
+      model.add(
+          keras.layers.Dense(
+              NUM_HIDDEN, input_dim=INPUT_DIM, activation='relu'))
+      model.add(keras.layers.Dense(NUM_CLASSES, activation='softmax'))
+      model.compile(
+          loss='binary_crossentropy',
+          optimizer=adam.AdamOptimizer(0.01),
+          metrics=['accuracy'])
+
+      cbks = [keras.callbacks.TensorBoard(log_dir=temp_dir)]
+
+      model.fit(
+          x_train,
+          y_train,
+          batch_size=BATCH_SIZE,
+          validation_data=(x_test, y_test),
+          callbacks=cbks,
+          epochs=2,
+          verbose=0)
+
+      self.assertTrue(os.path.exists(temp_dir))
+
   def test_RemoteMonitorWithJsonPayload(self):
     if requests is None:
       self.skipTest('`requests` required to run this test')