529 lines
16 KiB
Python
529 lines
16 KiB
Python
# Copyright 2017 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 helper routine module.
|
|
|
|
This module is a trove of succinct generic helper routines that don't
|
|
pull in any heavyweight dependencies aside from TensorFlow.
|
|
"""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import locale
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import numpy as np
|
|
import six
|
|
import tensorflow as tf
|
|
|
|
|
|
def setup_logging(streams=(sys.stderr,)):
|
|
"""Configures Python logging the way the TensorBoard team likes it.
|
|
|
|
This should be called exactly once at the beginning of main().
|
|
|
|
Args:
|
|
streams: An iterable of open files. Logs are written to each.
|
|
|
|
:type streams: tuple[file]
|
|
"""
|
|
# NOTE: Adding a level parameter to this method would be a bad idea
|
|
# because Python and ABSL disagree on the level numbers.
|
|
locale.setlocale(locale.LC_ALL, '')
|
|
tf.logging.set_verbosity(tf.logging.WARN)
|
|
# TODO(jart): Make the default TensorFlow logger behavior great again.
|
|
logging.currentframe = _hack_the_main_frame
|
|
handlers = [LogHandler(s) for s in streams]
|
|
formatter = LogFormatter()
|
|
for handler in handlers:
|
|
handler.setFormatter(formatter)
|
|
tensorflow_logger = logging.getLogger('tensorflow')
|
|
tensorflow_logger.handlers = handlers
|
|
tensorboard_logger = logging.getLogger('tensorboard')
|
|
tensorboard_logger.handlers = handlers
|
|
werkzeug_logger = logging.getLogger('werkzeug')
|
|
werkzeug_logger.setLevel(logging.WARNING)
|
|
werkzeug_logger.handlers = handlers
|
|
|
|
|
|
def closeable(class_):
|
|
"""Makes a class with a close method able to be a context manager.
|
|
|
|
This decorator is a great way to avoid having to choose between the
|
|
boilerplate of __enter__ and __exit__ methods, versus the boilerplate
|
|
of using contextlib.closing on every with statement.
|
|
|
|
Args:
|
|
class_: The class being decorated.
|
|
|
|
Raises:
|
|
ValueError: If class didn't have a close method, or already
|
|
implements __enter__ or __exit__.
|
|
"""
|
|
if 'close' not in class_.__dict__:
|
|
# coffee is for closers
|
|
raise ValueError('Class does not define a close() method: %s' % class_)
|
|
if '__enter__' in class_.__dict__ or '__exit__' in class_.__dict__:
|
|
raise ValueError('Class already defines __enter__ or __exit__: ' + class_)
|
|
class_.__enter__ = lambda self: self
|
|
class_.__exit__ = lambda self, t, v, b: self.close() and None
|
|
return class_
|
|
|
|
|
|
def close_all(resources):
|
|
"""Safely closes multiple resources.
|
|
|
|
The close method on all resources is guaranteed to be called. If
|
|
multiple close methods throw exceptions, then the first will be
|
|
raised and the rest will be logged.
|
|
|
|
Args:
|
|
resources: An iterable of object instances whose classes implement
|
|
the close method.
|
|
|
|
Raises:
|
|
Exception: To rethrow the last exception raised by a close method.
|
|
"""
|
|
exc_info = None
|
|
for resource in resources:
|
|
try:
|
|
resource.close()
|
|
except Exception as e: # pylint: disable=broad-except
|
|
if exc_info is not None:
|
|
tf.logging.error('Suppressing close(%s) failure: %s', resource, e,
|
|
exc_info=exc_info)
|
|
exc_info = sys.exc_info()
|
|
if exc_info is not None:
|
|
six.reraise(*exc_info)
|
|
|
|
|
|
def guarded_by(field):
|
|
"""Indicates method should be called from within a lock.
|
|
|
|
This decorator is purely for documentation purposes. It has the same
|
|
semantics as Java's @GuardedBy annotation.
|
|
|
|
Args:
|
|
field: The string name of the lock field, e.g. "_lock".
|
|
"""
|
|
del field
|
|
return lambda method: method
|
|
|
|
|
|
class Retrier(object):
|
|
"""Helper class for retrying things with exponential back-off."""
|
|
|
|
DELAY = 0.1
|
|
|
|
def __init__(self, is_transient, max_attempts=8, sleep=time.sleep):
|
|
"""Creates new instance.
|
|
|
|
:type is_transient: (Exception) -> bool
|
|
:type max_attempts: int
|
|
:type sleep: (float) -> None
|
|
"""
|
|
self._is_transient = is_transient
|
|
self._max_attempts = max_attempts
|
|
self._sleep = sleep
|
|
|
|
def run(self, callback):
|
|
"""Invokes callback, retrying on transient exceptions.
|
|
|
|
After the first failure, we wait 100ms, and then double with each
|
|
subsequent failed attempt. The default max attempts is 8 which
|
|
equates to about thirty seconds of sleeping total.
|
|
|
|
:type callback: () -> T
|
|
:rtype: T
|
|
"""
|
|
failures = 0
|
|
while True:
|
|
try:
|
|
return callback()
|
|
except Exception as e: # pylint: disable=broad-except
|
|
failures += 1
|
|
if failures == self._max_attempts or not self._is_transient(e):
|
|
raise
|
|
tf.logging.warn('Retrying on transient %s', e)
|
|
self._sleep(2 ** (failures - 1) * Retrier.DELAY)
|
|
|
|
|
|
class LogFormatter(logging.Formatter):
|
|
"""Google style log formatter.
|
|
|
|
The format is in essence the following:
|
|
|
|
[DIWEF]mmdd hh:mm:ss.uuuuuu thread_name file:line] msg
|
|
|
|
This class is meant to be used with LogHandler.
|
|
"""
|
|
|
|
DATE_FORMAT = '%m%d %H:%M:%S'
|
|
LOG_FORMAT = ('%(levelname)s%(asctime)s %(threadName)s '
|
|
'%(filename)s:%(lineno)d] %(message)s')
|
|
|
|
LEVEL_NAMES = {
|
|
logging.FATAL: 'F',
|
|
logging.ERROR: 'E',
|
|
logging.WARN: 'W',
|
|
logging.INFO: 'I',
|
|
logging.DEBUG: 'D',
|
|
}
|
|
|
|
def __init__(self):
|
|
"""Creates new instance."""
|
|
super(LogFormatter, self).__init__(LogFormatter.LOG_FORMAT,
|
|
LogFormatter.DATE_FORMAT)
|
|
|
|
def format(self, record):
|
|
"""Formats the log record.
|
|
|
|
:type record: logging.LogRecord
|
|
:rtype: str
|
|
"""
|
|
record.levelname = LogFormatter.LEVEL_NAMES[record.levelno]
|
|
return super(LogFormatter, self).format(record)
|
|
|
|
def formatTime(self, record, datefmt=None):
|
|
"""Return creation time of the specified LogRecord as formatted text.
|
|
|
|
This override adds microseconds.
|
|
|
|
:type record: logging.LogRecord
|
|
:rtype: str
|
|
"""
|
|
return (super(LogFormatter, self).formatTime(record, datefmt) +
|
|
'.%06d' % (record.created * 1e6 % 1e6))
|
|
|
|
|
|
class Ansi(object):
|
|
"""ANSI terminal codes container."""
|
|
|
|
ESCAPE = '\x1b['
|
|
ESCAPE_PATTERN = re.compile(re.escape(ESCAPE) + r'\??(?:\d+)(?:;\d+)*[mlh]')
|
|
RESET = ESCAPE + '0m'
|
|
BOLD = ESCAPE + '1m'
|
|
FLIP = ESCAPE + '7m'
|
|
RED = ESCAPE + '31m'
|
|
YELLOW = ESCAPE + '33m'
|
|
MAGENTA = ESCAPE + '35m'
|
|
CURSOR_HIDE = ESCAPE + '?25l'
|
|
CURSOR_SHOW = ESCAPE + '?25h'
|
|
|
|
|
|
class LogHandler(logging.StreamHandler):
|
|
"""Log handler that supports ANSI colors and ephemeral records.
|
|
|
|
Colors are applied on a line-by-line basis to non-INFO records. The
|
|
goal is to help the user visually distinguish meaningful information,
|
|
even when logging is verbose.
|
|
|
|
This handler will also strip ANSI color codes from emitted log
|
|
records automatically when the output stream is not a terminal.
|
|
|
|
Ephemeral log records are only emitted to a teletype emulator, only
|
|
display on the final row, and get overwritten as soon as another
|
|
ephemeral record is outputted. Ephemeral records are also sticky. If
|
|
a normal record is written then the previous ephemeral record is
|
|
restored right beneath it. When an ephemeral record with an empty
|
|
message is emitted, then the last ephemeral record turns into a
|
|
normal record and is allowed to spool.
|
|
|
|
This class is thread safe.
|
|
"""
|
|
|
|
EPHEMERAL = '.ephemeral' # Name suffix for ephemeral loggers.
|
|
|
|
COLORS = {
|
|
logging.FATAL: Ansi.BOLD + Ansi.RED,
|
|
logging.ERROR: Ansi.RED,
|
|
logging.WARN: Ansi.YELLOW,
|
|
logging.INFO: '',
|
|
logging.DEBUG: Ansi.MAGENTA,
|
|
}
|
|
|
|
def __init__(self, stream, type_='detect'):
|
|
"""Creates new instance.
|
|
|
|
Args:
|
|
stream: A file-like object.
|
|
type_: If "detect", will call stream.isatty() and perform system
|
|
checks to determine if it's safe to output ANSI terminal
|
|
codes. If type is "ansi" then this forces the use of ANSI
|
|
terminal codes.
|
|
|
|
Raises:
|
|
ValueError: If type is not "detect" or "ansi".
|
|
"""
|
|
if type_ not in ('detect', 'ansi'):
|
|
raise ValueError('type should be detect or ansi')
|
|
super(LogHandler, self).__init__(stream)
|
|
self._stream = stream
|
|
self._disable_flush = False
|
|
self._is_tty = (type_ == 'ansi' or
|
|
(hasattr(stream, 'isatty') and
|
|
stream.isatty() and
|
|
os.name != 'nt'))
|
|
self._ephemeral = ''
|
|
|
|
def emit(self, record):
|
|
"""Emits a log record.
|
|
|
|
:type record: logging.LogRecord
|
|
"""
|
|
self.acquire()
|
|
try:
|
|
is_ephemeral = record.name.endswith(LogHandler.EPHEMERAL)
|
|
color = LogHandler.COLORS.get(record.levelno)
|
|
if is_ephemeral:
|
|
if self._is_tty:
|
|
ephemeral = record.getMessage()
|
|
if ephemeral:
|
|
if color:
|
|
ephemeral = color + ephemeral + Ansi.RESET
|
|
self._clear_line()
|
|
self._stream.write(ephemeral)
|
|
else:
|
|
if self._ephemeral:
|
|
self._stream.write('\n')
|
|
self._ephemeral = ephemeral
|
|
else:
|
|
self._clear_line()
|
|
if self._is_tty and color:
|
|
self._stream.write(color)
|
|
self._disable_flush = True # prevent double flush
|
|
super(LogHandler, self).emit(record)
|
|
self._disable_flush = False
|
|
if self._is_tty and color:
|
|
self._stream.write(Ansi.RESET)
|
|
if self._ephemeral:
|
|
self._stream.write(self._ephemeral)
|
|
self.flush()
|
|
finally:
|
|
self._disable_flush = False
|
|
self.release()
|
|
|
|
def format(self, record):
|
|
"""Turns a log record into a string.
|
|
|
|
:type record: logging.LogRecord
|
|
:rtype: str
|
|
"""
|
|
message = super(LogHandler, self).format(record)
|
|
if not self._is_tty:
|
|
message = Ansi.ESCAPE_PATTERN.sub('', message)
|
|
return message
|
|
|
|
def flush(self):
|
|
"""Flushes output stream."""
|
|
self.acquire()
|
|
try:
|
|
if not self._disable_flush:
|
|
super(LogHandler, self).flush()
|
|
finally:
|
|
self.release()
|
|
|
|
def _clear_line(self):
|
|
if self._is_tty and self._ephemeral:
|
|
# We're counting columns in the terminal, not bytes. So we don't
|
|
# want to take UTF-8 or color codes into consideration.
|
|
text = Ansi.ESCAPE_PATTERN.sub('', tf.compat.as_text(self._ephemeral))
|
|
self._stream.write('\r' + ' ' * len(text) + '\r')
|
|
|
|
|
|
def _hack_the_main_frame():
|
|
"""Returns caller frame and skips over tf_logging.
|
|
|
|
This works around a bug in TensorFlow's open source logging module
|
|
where the Python logging module attributes log entries to the
|
|
delegate functions in tf_logging.py.
|
|
"""
|
|
if hasattr(sys, '_getframe'):
|
|
frame = sys._getframe(3)
|
|
else:
|
|
try:
|
|
raise Exception
|
|
except Exception: # pylint: disable=broad-except
|
|
frame = sys.exc_info()[2].tb_frame.f_back
|
|
if (frame is not None and
|
|
hasattr(frame.f_back, 'f_code') and
|
|
'tf_logging.py' in frame.f_back.f_code.co_filename):
|
|
return frame.f_back
|
|
return frame
|
|
|
|
|
|
class PersistentOpEvaluator(object):
|
|
"""Evaluate a fixed TensorFlow graph repeatedly, safely, efficiently.
|
|
|
|
Extend this class to create a particular kind of op evaluator, like an
|
|
image encoder. In `initialize_graph`, create an appropriate TensorFlow
|
|
graph with placeholder inputs. In `run`, evaluate this graph and
|
|
return its result. This class will manage a singleton graph and
|
|
session to preserve memory usage, and will ensure that this graph and
|
|
session do not interfere with other concurrent sessions.
|
|
|
|
A subclass of this class offers a threadsafe, highly parallel Python
|
|
entry point for evaluating a particular TensorFlow graph.
|
|
|
|
Example usage:
|
|
|
|
class FluxCapacitanceEvaluator(PersistentOpEvaluator):
|
|
\"\"\"Compute the flux capacitance required for a system.
|
|
|
|
Arguments:
|
|
x: Available power input, as a `float`, in jigawatts.
|
|
|
|
Returns:
|
|
A `float`, in nanofarads.
|
|
\"\"\"
|
|
|
|
def initialize_graph(self):
|
|
self._placeholder = tf.placeholder(some_dtype)
|
|
self._op = some_op(self._placeholder)
|
|
|
|
def run(self, x):
|
|
return self._op.eval(feed_dict: {self._placeholder: x})
|
|
|
|
evaluate_flux_capacitance = FluxCapacitanceEvaluator()
|
|
|
|
for x in xs:
|
|
evaluate_flux_capacitance(x)
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(PersistentOpEvaluator, self).__init__()
|
|
self._session = None
|
|
self._initialization_lock = threading.Lock()
|
|
|
|
def _lazily_initialize(self):
|
|
"""Initialize the graph and session, if this has not yet been done."""
|
|
with self._initialization_lock:
|
|
if self._session:
|
|
return
|
|
graph = tf.Graph()
|
|
with graph.as_default():
|
|
self.initialize_graph()
|
|
# Don't reserve GPU because libpng can't run on GPU.
|
|
config = tf.ConfigProto(device_count={'GPU': 0})
|
|
self._session = tf.Session(graph=graph, config=config)
|
|
|
|
def initialize_graph(self):
|
|
"""Create the TensorFlow graph needed to compute this operation.
|
|
|
|
This should write ops to the default graph and return `None`.
|
|
"""
|
|
raise NotImplementedError('Subclasses must implement "initialize_graph".')
|
|
|
|
def run(self, *args, **kwargs):
|
|
"""Evaluate the ops with the given input.
|
|
|
|
When this function is called, the default session will have the
|
|
graph defined by a previous call to `initialize_graph`. This
|
|
function should evaluate any ops necessary to compute the result of
|
|
the query for the given *args and **kwargs, likely returning the
|
|
result of a call to `some_op.eval(...)`.
|
|
"""
|
|
raise NotImplementedError('Subclasses must implement "run".')
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
self._lazily_initialize()
|
|
with self._session.as_default():
|
|
return self.run(*args, **kwargs)
|
|
|
|
|
|
class _TensorFlowPngEncoder(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(_TensorFlowPngEncoder, self).__init__()
|
|
self._image_placeholder = None
|
|
self._encode_op = None
|
|
|
|
def initialize_graph(self):
|
|
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(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(_TensorFlowWavEncoder, self).__init__()
|
|
self._audio_placeholder = None
|
|
self._samples_per_second_placeholder = None
|
|
self._encode_op = None
|
|
|
|
def initialize_graph(self):
|
|
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.contrib.ffmpeg.encode_audio(
|
|
self._audio_placeholder,
|
|
file_format='wav',
|
|
samples_per_second=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()
|