tensorbored 2.21.0rc1769983804__py3-none-any.whl

tensorbored/summary/writer/event_file_writer.py

@@ -0,0 +1,291 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Writes events to disk in a logdir."""
+
+import os
+import queue
+import socket
+import threading
+import time
+
+from tensorbored.compat import tf
+from tensorbored.compat.proto import event_pb2
+from tensorbored.summary.writer.record_writer import RecordWriter
+
+
+class AtomicCounter:
+    def __init__(self, initial_value):
+        self._value = initial_value
+        self._lock = threading.Lock()
+
+    def get(self):
+        with self._lock:
+            try:
+                return self._value
+            finally:
+                self._value += 1
+
+
+_global_uid = AtomicCounter(0)
+
+
+class EventFileWriter:
+    """Writes `Event` protocol buffers to an event file.
+
+    The `EventFileWriter` class creates an event file in the specified
+    directory, and asynchronously writes Event protocol buffers to the
+    file. The Event file is encoded using the tfrecord format, which is
+    similar to RecordIO.
+    """
+
+    def __init__(
+        self, logdir, max_queue_size=10, flush_secs=120, filename_suffix=""
+    ):
+        """Creates a `EventFileWriter` and an event file to write to.
+
+        On construction the summary writer creates a new event file in `logdir`.
+        This event file will contain `Event` protocol buffers, which are written to
+        disk via the add_event method.
+        The other arguments to the constructor control the asynchronous writes to
+        the event file:
+
+        Args:
+          logdir: A string. Directory where event file will be written.
+          max_queue_size: Integer. Size of the queue for pending events and summaries.
+          flush_secs: Number. How often, in seconds, to flush the
+            pending events and summaries to disk.
+        """
+        self._logdir = logdir
+        tf.io.gfile.makedirs(logdir)
+        self._file_name = (
+            os.path.join(
+                logdir,
+                "events.out.tfevents.%010d.%s.%s.%s"
+                % (
+                    time.time(),
+                    socket.gethostname(),
+                    os.getpid(),
+                    _global_uid.get(),
+                ),
+            )
+            + filename_suffix
+        )  # noqa E128
+        self._general_file_writer = tf.io.gfile.GFile(self._file_name, "wb")
+        self._async_writer = _AsyncWriter(
+            RecordWriter(self._general_file_writer), max_queue_size, flush_secs
+        )
+
+        # Initialize an event instance.
+        _event = event_pb2.Event(
+            wall_time=time.time(),
+            file_version="brain.Event:2",
+            source_metadata=event_pb2.SourceMetadata(
+                writer="tensorbored.summary.writer.event_file_writer"
+            ),
+        )
+        self.add_event(_event)
+        self.flush()
+
+    def get_logdir(self):
+        """Returns the directory where event file will be written."""
+        return self._logdir
+
+    def add_event(self, event):
+        """Adds an event to the event file.
+
+        Args:
+          event: An `Event` protocol buffer.
+        """
+        if not isinstance(event, event_pb2.Event):
+            raise TypeError(
+                "Expected an event_pb2.Event proto, "
+                " but got %s" % type(event)
+            )
+        self._async_writer.write(event.SerializeToString())
+
+    def flush(self):
+        """Flushes the event file to disk.
+
+        Call this method to make sure that all pending events have been
+        written to disk.
+        """
+        self._async_writer.flush()
+
+    def close(self):
+        """Performs a final flush of the event file to disk, stops the
+        write/flush worker and closes the file.
+
+        Call this method when you do not need the summary writer
+        anymore.
+        """
+        self._async_writer.close()
+
+
+class _AsyncWriter:
+    """Writes bytes to a file."""
+
+    def __init__(self, record_writer, max_queue_size=20, flush_secs=120):
+        """Writes bytes to a file asynchronously. An instance of this class
+        holds a queue to keep the incoming data temporarily. Data passed to the
+        `write` function will be put to the queue and the function returns
+        immediately. This class also maintains a thread to write data in the
+        queue to disk. The first initialization parameter is an instance of
+        `tensorboard.summary.record_writer` which computes the CRC checksum and
+        then write the combined result to the disk. So we use an async approach
+        to improve performance.
+
+        Args:
+            record_writer: A RecordWriter instance
+            max_queue_size: Integer. Size of the queue for pending bytestrings.
+            flush_secs: Number. How often, in seconds, to flush the
+                pending bytestrings to disk.
+        """
+        self._writer = record_writer
+        self._closed = False
+        self._byte_queue = queue.Queue(max_queue_size)
+        self._worker = _AsyncWriterThread(
+            self._byte_queue, self._writer, flush_secs
+        )
+        self._lock = threading.Lock()
+        self._worker.start()
+
+    def write(self, bytestring):
+        """Enqueue the given bytes to be written asychronously."""
+        with self._lock:
+            # Status of the worker should be checked under the lock to avoid
+            # multiple threads passing the check and then switching just before
+            # blocking on putting to the queue which might result in a deadlock.
+            self._check_worker_status()
+            if self._closed:
+                raise IOError("Writer is closed")
+            self._byte_queue.put(bytestring)
+            # Check the status again in case the background worker thread has
+            # failed in the meantime to avoid waiting until the next call to
+            # surface the error.
+            self._check_worker_status()
+
+    def flush(self):
+        """Write all the enqueued bytestring before this flush call to disk.
+
+        Block until all the above bytestring are written.
+        """
+        with self._lock:
+            self._check_worker_status()
+            if self._closed:
+                raise IOError("Writer is closed")
+            self._byte_queue.join()
+            self._writer.flush()
+            # Check the status again in case the background worker thread has
+            # failed in the meantime to avoid waiting until the next call to
+            # surface the error.
+            self._check_worker_status()
+
+    def close(self):
+        """Closes the underlying writer, flushing any pending writes first."""
+        if not self._closed:
+            with self._lock:
+                if not self._closed:
+                    self._closed = True
+                    self._worker.stop()
+                    self._writer.flush()
+                    self._writer.close()
+
+    def _check_worker_status(self):
+        """Makes sure the worker thread is still running and raises exception
+        thrown in the worker thread otherwise.
+        """
+        exception = self._worker.exception
+        if exception is not None:
+            raise exception
+
+
+class _AsyncWriterThread(threading.Thread):
+    """Thread that processes asynchronous writes for _AsyncWriter."""
+
+    def __init__(self, queue, record_writer, flush_secs):
+        """Creates an _AsyncWriterThread.
+
+        Args:
+          queue: A Queue from which to dequeue data.
+          record_writer: An instance of record_writer writer.
+          flush_secs: How often, in seconds, to flush the
+            pending file to disk.
+        """
+        threading.Thread.__init__(self)
+        self.daemon = True
+        self.exception = None
+        self._queue = queue
+        self._record_writer = record_writer
+        self._flush_secs = flush_secs
+        # The first data will be flushed immediately.
+        self._next_flush_time = 0
+        self._has_pending_data = False
+        self._shutdown_signal = object()
+
+    def stop(self):
+        self._queue.put(self._shutdown_signal)
+        self.join()
+
+    def run(self):
+        try:
+            self._run()
+        except Exception as ex:
+            self.exception = ex
+            try:
+                # In case there's a thread blocked on putting an item into the
+                # queue or a thread blocked on flushing, pop all items from the
+                # queue to let the foreground thread proceed.
+                while True:
+                    self._queue.get(False)
+                    self._queue.task_done()
+            except queue.Empty:
+                pass
+            raise
+
+    def _run(self):
+        # Here wait on the queue until an data appears, or till the next
+        # time to flush the writer, whichever is earlier. If we have an
+        # data, write it. If not, an empty queue exception will be raised
+        # and we can proceed to flush the writer.
+        while True:
+            now = time.time()
+            queue_wait_duration = self._next_flush_time - now
+            data = None
+            try:
+                if queue_wait_duration > 0:
+                    data = self._queue.get(True, queue_wait_duration)
+                else:
+                    data = self._queue.get(False)
+
+                if data is self._shutdown_signal:
+                    return
+                self._record_writer.write(data)
+                self._has_pending_data = True
+            except queue.Empty:
+                pass
+            finally:
+                if data:
+                    self._queue.task_done()
+
+            now = time.time()
+            if now > self._next_flush_time:
+                if self._has_pending_data:
+                    # Small optimization - if there are no pending data,
+                    # there's no need to flush, since each flush can be
+                    # expensive (e.g. uploading a new file to a server).
+                    self._record_writer.flush()
+                    self._has_pending_data = False
+                # Do it again in flush_secs.
+                self._next_flush_time = now + self._flush_secs

tensorbored/summary/writer/record_writer.py

@@ -0,0 +1,50 @@
+# Copyright 2019 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.
+# ==============================================================================
+
+import struct
+from tensorbored.compat.tensorflow_stub.pywrap_tensorflow import masked_crc32c
+
+
+class RecordWriter:
+    """Write encoded protobuf to a file with packing defined in tensorflow."""
+
+    def __init__(self, writer):
+        """Open a file to keep the tensorboard records.
+
+        Args:
+        writer: A file-like object that implements `write`, `flush` and `close`.
+        """
+        self._writer = writer
+
+    # Format of a single record: (little-endian)
+    # uint64    length
+    # uint32    masked crc of length
+    # byte      data[length]
+    # uint32    masked crc of data
+    def write(self, data):
+        header = struct.pack("<Q", len(data))
+        header_crc = struct.pack("<I", masked_crc32c(header))
+        footer_crc = struct.pack("<I", masked_crc32c(data))
+        self._writer.write(header + header_crc + data + footer_crc)
+
+    def flush(self):
+        self._writer.flush()
+
+    def close(self):
+        self._writer.close()
+
+    @property
+    def closed(self):
+        return self._writer.closed

tensorbored/util/encoder.py

@@ -0,0 +1,116 @@
+# Copyright 2018 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.
+
+"""TensorBoard encoder helper module.
+
+Encoder depends on TensorFlow.
+"""
+
+import numpy as np
+
+from tensorbored.util import op_evaluator
+
+
+class _TensorFlowPngEncoder(op_evaluator.PersistentOpEvaluator):
+    """Encode an image to PNG.
+
+    This function is thread-safe, and has high performance when run in
+    parallel. See `encode_png_benchmark.py` for details.
+
+    Arguments:
+      image: A numpy array of shape `[height, width, channels]`, where
+        `channels` is 1, 3, or 4, and of dtype uint8.
+
+    Returns:
+      A bytestring with PNG-encoded data.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._image_placeholder = None
+        self._encode_op = None
+
+    def initialize_graph(self):
+        # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
+        import tensorflow.compat.v1 as tf
+
+        self._image_placeholder = tf.placeholder(
+            dtype=tf.uint8, name="image_to_encode"
+        )
+        self._encode_op = tf.image.encode_png(self._image_placeholder)
+
+    def run(self, image):  # pylint: disable=arguments-differ
+        if not isinstance(image, np.ndarray):
+            raise ValueError("'image' must be a numpy array: %r" % image)
+        if image.dtype != np.uint8:
+            raise ValueError(
+                "'image' dtype must be uint8, but is %r" % image.dtype
+            )
+        return self._encode_op.eval(feed_dict={self._image_placeholder: image})
+
+
+encode_png = _TensorFlowPngEncoder()
+
+
+class _TensorFlowWavEncoder(op_evaluator.PersistentOpEvaluator):
+    """Encode an audio clip to WAV.
+
+    This function is thread-safe and exhibits good parallel performance.
+
+    Arguments:
+      audio: A numpy array of shape `[samples, channels]`.
+      samples_per_second: A positive `int`, in Hz.
+
+    Returns:
+      A bytestring with WAV-encoded data.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._audio_placeholder = None
+        self._samples_per_second_placeholder = None
+        self._encode_op = None
+
+    def initialize_graph(self):
+        # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
+        import tensorflow.compat.v1 as tf
+
+        self._audio_placeholder = tf.placeholder(
+            dtype=tf.float32, name="image_to_encode"
+        )
+        self._samples_per_second_placeholder = tf.placeholder(
+            dtype=tf.int32, name="samples_per_second"
+        )
+        self._encode_op = tf.audio.encode_wav(
+            self._audio_placeholder,
+            sample_rate=self._samples_per_second_placeholder,
+        )
+
+    def run(
+        self, audio, samples_per_second
+    ):  # pylint: disable=arguments-differ
+        if not isinstance(audio, np.ndarray):
+            raise ValueError("'audio' must be a numpy array: %r" % audio)
+        if not isinstance(samples_per_second, int):
+            raise ValueError(
+                "'samples_per_second' must be an int: %r" % samples_per_second
+            )
+        feed_dict = {
+            self._audio_placeholder: audio,
+            self._samples_per_second_placeholder: samples_per_second,
+        }
+        return self._encode_op.eval(feed_dict=feed_dict)
+
+
+encode_wav = _TensorFlowWavEncoder()