• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""TensorBoard Plugin asset abstract class.
16
17TensorBoard plugins may need to provide arbitrary assets, such as
18configuration information for specific outputs, or vocabulary files, or sprite
19images, etc.
20
21This module contains methods that allow plugin assets to be specified at graph
22construction time. Plugin authors define a PluginAsset which is treated as a
23singleton on a per-graph basis. The PluginAsset has an assets method which
24returns a dictionary of asset contents. The tf.compat.v1.summary.FileWriter
25(or any other Summary writer) will serialize these assets in such a way that
26TensorBoard can retrieve them.
27"""
28
29from __future__ import absolute_import
30from __future__ import division
31from __future__ import print_function
32
33import abc
34
35import six
36
37from tensorflow.python.framework import ops
38
39_PLUGIN_ASSET_PREFIX = "__tensorboard_plugin_asset__"
40
41
42def get_plugin_asset(plugin_asset_cls, graph=None):
43  """Acquire singleton PluginAsset instance from a graph.
44
45  PluginAssets are always singletons, and are stored in tf Graph collections.
46  This way, they can be defined anywhere the graph is being constructed, and
47  if the same plugin is configured at many different points, the user can always
48  modify the same instance.
49
50  Args:
51    plugin_asset_cls: The PluginAsset class
52    graph: (optional) The graph to retrieve the instance from. If not specified,
53      the default graph is used.
54
55  Returns:
56    An instance of the plugin_asset_class
57
58  Raises:
59    ValueError: If we have a plugin name collision, or if we unexpectedly find
60      the wrong number of items in a collection.
61  """
62  if graph is None:
63    graph = ops.get_default_graph()
64  if not plugin_asset_cls.plugin_name:
65    raise ValueError("Class %s has no plugin_name" % plugin_asset_cls.__name__)
66
67  name = _PLUGIN_ASSET_PREFIX + plugin_asset_cls.plugin_name
68  container = graph.get_collection(name)
69  if container:
70    if len(container) != 1:
71      raise ValueError("Collection for %s had %d items, expected 1" %
72                       (name, len(container)))
73    instance = container[0]
74    if not isinstance(instance, plugin_asset_cls):
75      raise ValueError("Plugin name collision between classes %s and %s" %
76                       (plugin_asset_cls.__name__, instance.__class__.__name__))
77  else:
78    instance = plugin_asset_cls()
79    graph.add_to_collection(name, instance)
80    graph.add_to_collection(_PLUGIN_ASSET_PREFIX, plugin_asset_cls.plugin_name)
81  return instance
82
83
84def get_all_plugin_assets(graph=None):
85  """Retrieve all PluginAssets stored in the graph collection.
86
87  Args:
88    graph: Optionally, the graph to get assets from. If unspecified, the default
89      graph is used.
90
91  Returns:
92    A list with all PluginAsset instances in the graph.
93
94  Raises:
95    ValueError: if we unexpectedly find a collection with the wrong number of
96      PluginAssets.
97
98  """
99  if graph is None:
100    graph = ops.get_default_graph()
101
102  out = []
103  for name in graph.get_collection(_PLUGIN_ASSET_PREFIX):
104    collection = graph.get_collection(_PLUGIN_ASSET_PREFIX + name)
105    if len(collection) != 1:
106      raise ValueError("Collection for %s had %d items, expected 1" %
107                       (name, len(collection)))
108    out.append(collection[0])
109  return out
110
111
112@six.add_metaclass(abc.ABCMeta)
113class PluginAsset(object):
114  """This abstract base class allows TensorBoard to serialize assets to disk.
115
116  Plugin authors are expected to extend the PluginAsset class, so that it:
117  - has a unique plugin_name
118  - provides an assets method that returns an {asset_name: asset_contents}
119    dictionary. For now, asset_contents are strings, although we may add
120    StringIO support later.
121
122  LifeCycle of a PluginAsset instance:
123  - It is constructed when get_plugin_asset is called on the class for
124    the first time.
125  - It is configured by code that follows the calls to get_plugin_asset
126  - When the containing graph is serialized by the
127    tf.compat.v1.summary.FileWriter, the writer calls assets and the
128    PluginAsset instance provides its contents to be written to disk.
129  """
130
131  plugin_name = None
132
133  @abc.abstractmethod
134  def assets(self):
135    """Provide all of the assets contained by the PluginAsset instance.
136
137    The assets method should return a dictionary structured as
138    {asset_name: asset_contents}. asset_contents is a string.
139
140    This method will be called by the tf.compat.v1.summary.FileWriter when it
141    is time to write the assets out to disk.
142    """
143    raise NotImplementedError()
144