Use the new SummaryMetadata field of the Value proto to pass per-plugin information.

This will let TensorBoard figure out which plugins should handle which types of summaries.

PiperOrigin-RevId: 159202469

5 files changed, 111 insertions, 12 deletions

diff --git a/tensorflow/python/summary/summary.py b/tensorflow/python/summary/summary.py
index 7ff01a51f3..f3600793a6 100644
--- a/tensorflow/python/summary/summary.py
+++ b/tensorflow/python/summary/summary.py
@@ -20,6 +20,7 @@ See the @{$python/summary} guide.
 @@FileWriter
 @@FileWriterCache
 @@tensor_summary
+@@_tensor_summary_v2
 @@scalar
 @@histogram
 @@audio
@@ -28,6 +29,7 @@ See the @{$python/summary} guide.
 @@merge
 @@merge_all
 @@get_summary_description
+@@PluginAsset
 @@get_plugin_asset
 @@get_all_plugin_assets
 """
diff --git a/tensorflow/python/summary/text_summary.py b/tensorflow/python/summary/text_summary.py
index 52bc913b2a..2132dc6eb8 100644
--- a/tensorflow/python/summary/text_summary.py
+++ b/tensorflow/python/summary/text_summary.py
@@ -23,13 +23,26 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from collections import namedtuple
 import json
 
+from tensorflow.core.framework import summary_pb2
 from tensorflow.python.framework import dtypes
-from tensorflow.python.ops.summary_ops import tensor_summary
+from tensorflow.python.ops.summary_ops import _tensor_summary_v2
 from tensorflow.python.summary import plugin_asset
+from tensorflow.python.util import deprecation
 
+PLUGIN_NAME = "text"
 
+# Contains event-related data specific to the text plugin.
+_TextPluginData = namedtuple("_TextPluginData", [])
+
+
+@deprecation.deprecated_args(
+    "2017-06-13",
+    "collections is deprecated. Instead of using collections to associate "
+    "plugins to events, add a PluginData field to the SummaryMetadata of a "
+    "Value proto.", "collections")
 def text_summary(name, tensor, collections=None):
   """Summarizes textual data.
 
@@ -60,9 +73,16 @@ def text_summary(name, tensor, collections=None):
     raise ValueError("Expected tensor %s to have dtype string, got %s" %
                      (tensor.name, tensor.dtype))
 
-  t_summary = tensor_summary(name, tensor, collections=collections)
-  text_assets = plugin_asset.get_plugin_asset(TextSummaryPluginAsset)
-  text_assets.register_tensor(t_summary.op.name)
+  summary_metadata = summary_pb2.SummaryMetadata()
+  text_plugin_data = _TextPluginData()
+  data_dict = text_plugin_data._asdict()  # pylint: disable=protected-access
+  summary_metadata.plugin_data.add(
+      plugin_name=PLUGIN_NAME, content=json.dumps(data_dict))
+  t_summary = _tensor_summary_v2(
+      name=name,
+      tensor=tensor,
+      summary_metadata=summary_metadata,
+      collections=collections)
   return t_summary
 
 
diff --git a/tensorflow/python/summary/text_summary_test.py b/tensorflow/python/summary/text_summary_test.py
index 31009702ca..4d357918f6 100644
--- a/tensorflow/python/summary/text_summary_test.py
+++ b/tensorflow/python/summary/text_summary_test.py
@@ -17,7 +17,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.framework import ops as framework_ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import googletest
@@ -43,16 +42,11 @@ class TextPluginTest(test_util.TensorFlowTestCase):
       # The API accepts vectors.
       arr = array_ops.constant(["one", "two", "three"])
       summ = text_summary.text_summary("foo", arr)
-      self.assertEqual(summ.op.type, "TensorSummary")
+      self.assertEqual(summ.op.type, "TensorSummaryV2")
 
       # the API accepts scalars
       summ = text_summary.text_summary("foo", array_ops.constant("one"))
-      self.assertEqual(summ.op.type, "TensorSummary")
-
-  def testTextSummaryCollections(self):
-    text_summary.text_summary("bar", array_ops.constant("2"), collections=[])
-    summaries = framework_ops.get_collection(framework_ops.GraphKeys.SUMMARIES)
-    self.assertEqual(len(summaries), 0)
+      self.assertEqual(summ.op.type, "TensorSummaryV2")
 
 
 if __name__ == "__main__":
diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py
index 05f97fb284..8ce49d623d 100644
--- a/tensorflow/python/summary/writer/writer.py
+++ b/tensorflow/python/summary/writer/writer.py
@@ -86,6 +86,14 @@ class SummaryToEventTransformer(object):
           meta_graph.create_meta_graph_def(graph_def=graph_def or
                                            maybe_graph_as_def))
 
+    # This set contains tags of Summary Values that have been encountered
+    # already. The motivation here is that the SummaryWriter only keeps the
+    # metadata property (which is a SummaryMetadata proto) of the first Summary
+    # Value encountered for each tag. The SummaryWriter strips away the
+    # SummaryMetadata for all subsequent Summary Values with tags seen
+    # previously. This saves space.
+    self._seen_summary_tags = set()
+
   def add_summary(self, summary, global_step=None):
     """Adds a `Summary` protocol buffer to the event file.
 
@@ -108,6 +116,24 @@ class SummaryToEventTransformer(object):
       summ = summary_pb2.Summary()
       summ.ParseFromString(summary)
       summary = summ
+
+    # We strip metadata from values with tags that we have seen before in order
+    # to save space - we just store the metadata on the first value with a
+    # specific tag.
+    for value in summary.value:
+      if not value.metadata:
+        continue
+
+      if value.tag in self._seen_summary_tags:
+        # This tag has been encountered before. Strip the metadata.
+        value.ClearField("metadata")
+        continue
+
+      # We encounter a value with a tag we have not encountered previously. And
+      # it has metadata. Remember to strip metadata from future values with this
+      # tag string.
+      self._seen_summary_tags.add(value.tag)
+
     event = event_pb2.Event(summary=summary)
     self._add_event(event, global_step)
 
diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py
index 8c34eb82e3..3d27b11cb9 100644
--- a/tensorflow/python/summary/writer/writer_test.py
+++ b/tensorflow/python/summary/writer/writer_test.py
@@ -317,6 +317,63 @@ class SummaryWriterTestCase(test.TestCase):
     # We should be done.
     self.assertRaises(StopIteration, lambda: next(rr))
 
+  def testPluginMetadataStrippedFromSubsequentEvents(self):
+    test_dir = self._CleanTestDir("basics")
+    sw = writer.FileWriter(test_dir)
+
+    sw.add_session_log(event_pb2.SessionLog(status=SessionLog.START), 1)
+
+    # We add 2 summaries with the same tags. They both have metadata. The writer
+    # should strip the metadata from the second one.
+    value = summary_pb2.Summary.Value(tag="foo", simple_value=10.0)
+    value.metadata.plugin_data.add(plugin_name="bar", content="... content ...")
+    sw.add_summary(summary_pb2.Summary(value=[value]), 10)
+    value = summary_pb2.Summary.Value(tag="foo", simple_value=10.0)
+    value.metadata.plugin_data.add(plugin_name="bar", content="... content ...")
+    sw.add_summary(summary_pb2.Summary(value=[value]), 10)
+
+    sw.close()
+    rr = self._EventsReader(test_dir)
+
+    # The first event should list the file_version.
+    ev = next(rr)
+    self._assertRecent(ev.wall_time)
+    self.assertEquals("brain.Event:2", ev.file_version)
+
+    # The next event should be the START message.
+    ev = next(rr)
+    self._assertRecent(ev.wall_time)
+    self.assertEquals(1, ev.step)
+    self.assertEquals(SessionLog.START, ev.session_log.status)
+
+    # This is the first event with tag foo. It should contain SummaryMetadata.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value {
+        tag: "foo"
+        simple_value: 10.0
+        metadata {
+          plugin_data {
+            plugin_name: "bar"
+            content: "... content ..."
+          }
+        }
+      }
+      """, ev.summary)
+
+    # This is the second event with tag foo. It should lack SummaryMetadata
+    # because the file writer should have stripped it.
+    ev = next(rr)
+    self.assertProtoEquals("""
+      value {
+        tag: "foo"
+        simple_value: 10.0
+      }
+      """, ev.summary)
+
+    # We should be done.
+    self.assertRaises(StopIteration, lambda: next(rr))
+
   def testFileWriterWithSuffix(self):
     test_dir = self._CleanTestDir("test_suffix")
     sw = writer.FileWriter(test_dir, filename_suffix="_test_suffix")