125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# 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
|
|
#
|
|
# http://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.
|
|
# ==============================================================================
|
|
"""Generalized output options for writing tensor-formatted summary data."""
|
|
|
|
from tensorboard.compat.proto import event_pb2
|
|
from tensorboard.compat.proto import summary_pb2
|
|
from tensorboard.summary.writer import event_file_writer
|
|
from tensorboard.util import tensor_util
|
|
|
|
import abc
|
|
|
|
|
|
class Output(abc.ABC):
|
|
"""Interface for emitting tensor-formatted summary data.
|
|
|
|
Implementations of this interface can be passed to Writer to customize
|
|
how summary data is actually persisted (e.g. to disk, to memory, over
|
|
the network, etc.).
|
|
|
|
TODO(#4581): This API should be considered EXPERIMENTAL and subject to
|
|
backwards-incompatible changes without notice.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def emit_scalar(
|
|
self,
|
|
*,
|
|
plugin_name,
|
|
tag,
|
|
data,
|
|
step,
|
|
wall_time,
|
|
tag_metadata=None,
|
|
description=None,
|
|
):
|
|
"""Emits one scalar data point to this Output.
|
|
|
|
Args:
|
|
plugin_name: string name to uniquely identify the type of time series
|
|
(historically associated with a TensorBoard plugin).
|
|
tag: string tag used to uniquely identify this time series.
|
|
data: `np.float32` scalar value for this data point.
|
|
step: `np.int64` scalar step value for this data point.
|
|
wall_time: `float` seconds since the Unix epoch, representing the
|
|
real-world timestamp for this data point.
|
|
tag_metadata: optional bytes containing metadata for this entire time
|
|
series. This should be constant for a given tag; only the first
|
|
value encountered will be used.
|
|
description: optional string description for this entire time series.
|
|
This should be constant for a given tag; only the first value
|
|
encountered will be used.
|
|
"""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def flush(self):
|
|
"""Flushes any data that has been buffered."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def close(self):
|
|
"""Closes the Output and also flushes any buffered data."""
|
|
pass
|
|
|
|
|
|
class DirectoryOutput(Output):
|
|
"""Outputs summary data by writing event files to a log directory.
|
|
|
|
TODO(#4581): This API should be considered EXPERIMENTAL and subject to
|
|
backwards-incompatible changes without notice.
|
|
"""
|
|
|
|
def __init__(self, path):
|
|
"""Creates a `DirectoryOutput` for the given path."""
|
|
self._ev_writer = event_file_writer.EventFileWriter(path)
|
|
|
|
def emit_scalar(
|
|
self,
|
|
*,
|
|
plugin_name,
|
|
tag,
|
|
data,
|
|
step,
|
|
wall_time,
|
|
tag_metadata=None,
|
|
description=None,
|
|
):
|
|
"""See `Output`."""
|
|
# TODO(#4581): cache summary metadata to emit only once.
|
|
summary_metadata = summary_pb2.SummaryMetadata(
|
|
plugin_data=summary_pb2.SummaryMetadata.PluginData(
|
|
plugin_name=plugin_name, content=tag_metadata
|
|
),
|
|
summary_description=description,
|
|
data_class=summary_pb2.DataClass.DATA_CLASS_SCALAR,
|
|
)
|
|
tensor_proto = tensor_util.make_tensor_proto(data)
|
|
event = event_pb2.Event(wall_time=wall_time, step=step)
|
|
event.summary.value.add(
|
|
tag=tag, tensor=tensor_proto, metadata=summary_metadata
|
|
)
|
|
self._ev_writer.add_event(event)
|
|
|
|
def flush(self):
|
|
"""See `Output`."""
|
|
self._ev_writer.flush()
|
|
|
|
def close(self):
|
|
"""See `Output`."""
|
|
# No need to call flush first since EventFileWriter already
|
|
# will do this for us when we call close().
|
|
self._ev_writer.close()
|