Viewing file:  util.py (16.44 KB) -rw-r--r--
Select action/file-type:
 (+) |  (+) |  (+) | Code (+) | Session (+) |  (+) | SDB (+) |  (+) |  (+) |  (+) |  (+) |  (+) |
# 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()
|