308 lines
9.9 KiB
Python
308 lines
9.9 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.
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from collections import deque
|
|
from math import floor, sqrt
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
from tensorboard.plugins.beholder import im_util
|
|
from tensorboard.plugins.beholder.shared_config import SECTION_HEIGHT,\
|
|
IMAGE_WIDTH, DEFAULT_CONFIG, SECTION_INFO_FILENAME
|
|
from tensorboard.plugins.beholder.file_system_tools import write_pickle
|
|
|
|
MIN_SQUARE_SIZE = 3
|
|
|
|
|
|
class Visualizer(object):
|
|
|
|
def __init__(self, logdir):
|
|
self.logdir = logdir
|
|
self.sections_over_time = deque([], DEFAULT_CONFIG['window_size'])
|
|
self.config = dict(DEFAULT_CONFIG)
|
|
self.old_config = dict(DEFAULT_CONFIG)
|
|
|
|
|
|
def _reshape_conv_array(self, array, section_height, image_width):
|
|
'''Reshape a rank 4 array to be rank 2, where each column of block_width is
|
|
a filter, and each row of block height is an input channel. For example:
|
|
|
|
[[[[ 11, 21, 31, 41],
|
|
[ 51, 61, 71, 81],
|
|
[ 91, 101, 111, 121]],
|
|
[[ 12, 22, 32, 42],
|
|
[ 52, 62, 72, 82],
|
|
[ 92, 102, 112, 122]],
|
|
[[ 13, 23, 33, 43],
|
|
[ 53, 63, 73, 83],
|
|
[ 93, 103, 113, 123]]],
|
|
[[[ 14, 24, 34, 44],
|
|
[ 54, 64, 74, 84],
|
|
[ 94, 104, 114, 124]],
|
|
[[ 15, 25, 35, 45],
|
|
[ 55, 65, 75, 85],
|
|
[ 95, 105, 115, 125]],
|
|
[[ 16, 26, 36, 46],
|
|
[ 56, 66, 76, 86],
|
|
[ 96, 106, 116, 126]]],
|
|
[[[ 17, 27, 37, 47],
|
|
[ 57, 67, 77, 87],
|
|
[ 97, 107, 117, 127]],
|
|
[[ 18, 28, 38, 48],
|
|
[ 58, 68, 78, 88],
|
|
[ 98, 108, 118, 128]],
|
|
[[ 19, 29, 39, 49],
|
|
[ 59, 69, 79, 89],
|
|
[ 99, 109, 119, 129]]]]
|
|
|
|
should be reshaped to:
|
|
|
|
[[ 11, 12, 13, 21, 22, 23, 31, 32, 33, 41, 42, 43],
|
|
[ 14, 15, 16, 24, 25, 26, 34, 35, 36, 44, 45, 46],
|
|
[ 17, 18, 19, 27, 28, 29, 37, 38, 39, 47, 48, 49],
|
|
[ 51, 52, 53, 61, 62, 63, 71, 72, 73, 81, 82, 83],
|
|
[ 54, 55, 56, 64, 65, 66, 74, 75, 76, 84, 85, 86],
|
|
[ 57, 58, 59, 67, 68, 69, 77, 78, 79, 87, 88, 89],
|
|
[ 91, 92, 93, 101, 102, 103, 111, 112, 113, 121, 122, 123],
|
|
[ 94, 95, 96, 104, 105, 106, 114, 115, 116, 124, 125, 126],
|
|
[ 97, 98, 99, 107, 108, 109, 117, 118, 119, 127, 128, 129]]
|
|
'''
|
|
|
|
# E.g. [100, 24, 24, 10]: this shouldn't be reshaped like normal.
|
|
if array.shape[1] == array.shape[2] and array.shape[0] != array.shape[1]:
|
|
array = np.rollaxis(np.rollaxis(array, 2), 2)
|
|
|
|
block_height, block_width, in_channels = array.shape[:3]
|
|
rows = []
|
|
|
|
max_element_count = section_height * int(image_width / MIN_SQUARE_SIZE)
|
|
element_count = 0
|
|
|
|
for i in range(in_channels):
|
|
rows.append(array[:, :, i, :].reshape(block_height, -1, order='F'))
|
|
|
|
# This line should be left in this position. Gives it one extra row.
|
|
if element_count >= max_element_count and not self.config['show_all']:
|
|
break
|
|
|
|
element_count += block_height * in_channels * block_width
|
|
|
|
return np.vstack(rows)
|
|
|
|
|
|
def _reshape_irregular_array(self, array, section_height, image_width):
|
|
'''Reshapes arrays of ranks not in {1, 2, 4}
|
|
'''
|
|
section_area = section_height * image_width
|
|
flattened_array = np.ravel(array)
|
|
|
|
if not self.config['show_all']:
|
|
flattened_array = flattened_array[:int(section_area/MIN_SQUARE_SIZE)]
|
|
|
|
cell_count = np.prod(flattened_array.shape)
|
|
cell_area = section_area / cell_count
|
|
|
|
cell_side_length = max(1, floor(sqrt(cell_area)))
|
|
row_count = max(1, int(section_height / cell_side_length))
|
|
col_count = int(cell_count / row_count)
|
|
|
|
# Reshape the truncated array so that it has the same aspect ratio as
|
|
# the section.
|
|
|
|
# Truncate whatever remaining values there are that don't fit. Hopefully
|
|
# it doesn't matter that the last few (< section count) aren't there.
|
|
section = np.reshape(flattened_array[:row_count * col_count],
|
|
(row_count, col_count))
|
|
|
|
return section
|
|
|
|
|
|
def _determine_image_width(self, arrays, show_all):
|
|
final_width = IMAGE_WIDTH
|
|
|
|
if show_all:
|
|
for array in arrays:
|
|
rank = len(array.shape)
|
|
|
|
if rank == 1:
|
|
width = len(array)
|
|
elif rank == 2:
|
|
width = array.shape[1]
|
|
elif rank == 4:
|
|
width = array.shape[1] * array.shape[3]
|
|
else:
|
|
width = IMAGE_WIDTH
|
|
|
|
if width > final_width:
|
|
final_width = width
|
|
|
|
return final_width
|
|
|
|
|
|
def _determine_section_height(self, array, show_all):
|
|
rank = len(array.shape)
|
|
height = SECTION_HEIGHT
|
|
|
|
if show_all:
|
|
if rank == 1:
|
|
height = SECTION_HEIGHT
|
|
if rank == 2:
|
|
height = max(SECTION_HEIGHT, array.shape[0])
|
|
elif rank == 4:
|
|
height = max(SECTION_HEIGHT, array.shape[0] * array.shape[2])
|
|
else:
|
|
height = max(SECTION_HEIGHT, np.prod(array.shape) // IMAGE_WIDTH)
|
|
|
|
return height
|
|
|
|
|
|
def _arrays_to_sections(self, arrays):
|
|
'''
|
|
input: unprocessed numpy arrays.
|
|
returns: columns of the size that they will appear in the image, not scaled
|
|
for display. That needs to wait until after variance is computed.
|
|
'''
|
|
sections = []
|
|
sections_to_resize_later = {}
|
|
show_all = self.config['show_all']
|
|
image_width = self._determine_image_width(arrays, show_all)
|
|
|
|
for array_number, array in enumerate(arrays):
|
|
rank = len(array.shape)
|
|
section_height = self._determine_section_height(array, show_all)
|
|
|
|
if rank == 1:
|
|
section = np.atleast_2d(array)
|
|
elif rank == 2:
|
|
section = array
|
|
elif rank == 4:
|
|
section = self._reshape_conv_array(array, section_height, image_width)
|
|
else:
|
|
section = self._reshape_irregular_array(array,
|
|
section_height,
|
|
image_width)
|
|
# Only calculate variance for what we have to. In some cases (biases),
|
|
# the section is larger than the array, so we don't want to calculate
|
|
# variance for the same value over and over - better to resize later.
|
|
# About a 6-7x speedup for a big network with a big variance window.
|
|
section_size = section_height * image_width
|
|
array_size = np.prod(array.shape)
|
|
|
|
if section_size > array_size:
|
|
sections.append(section)
|
|
sections_to_resize_later[array_number] = section_height
|
|
else:
|
|
sections.append(im_util.resize(section, section_height, image_width))
|
|
|
|
self.sections_over_time.append(sections)
|
|
|
|
if self.config['mode'] == 'variance':
|
|
sections = self._sections_to_variance_sections(self.sections_over_time)
|
|
|
|
for array_number, height in sections_to_resize_later.items():
|
|
sections[array_number] = im_util.resize(sections[array_number],
|
|
height,
|
|
image_width)
|
|
return sections
|
|
|
|
|
|
def _sections_to_variance_sections(self, sections_over_time):
|
|
'''Computes the variance of corresponding sections over time.
|
|
|
|
Returns:
|
|
a list of np arrays.
|
|
'''
|
|
variance_sections = []
|
|
|
|
for i in range(len(sections_over_time[0])):
|
|
time_sections = [sections[i] for sections in sections_over_time]
|
|
variance = np.var(time_sections, axis=0)
|
|
variance_sections.append(variance)
|
|
|
|
return variance_sections
|
|
|
|
|
|
def _sections_to_image(self, sections):
|
|
padding_size = 5
|
|
|
|
sections = im_util.scale_sections(sections, self.config['scaling'])
|
|
|
|
final_stack = [sections[0]]
|
|
padding = np.zeros((padding_size, sections[0].shape[1]))
|
|
|
|
for section in sections[1:]:
|
|
final_stack.append(padding)
|
|
final_stack.append(section)
|
|
|
|
return np.vstack(final_stack).astype(np.uint8)
|
|
|
|
|
|
def _maybe_clear_deque(self):
|
|
'''Clears the deque if certain parts of the config have changed.'''
|
|
|
|
for config_item in ['values', 'mode', 'show_all']:
|
|
if self.config[config_item] != self.old_config[config_item]:
|
|
self.sections_over_time.clear()
|
|
break
|
|
|
|
self.old_config = self.config
|
|
|
|
window_size = self.config['window_size']
|
|
if window_size != self.sections_over_time.maxlen:
|
|
self.sections_over_time = deque(self.sections_over_time, window_size)
|
|
|
|
|
|
def _save_section_info(self, arrays, sections):
|
|
infos = []
|
|
|
|
if self.config['values'] == 'trainable_variables':
|
|
names = [x.name for x in tf.trainable_variables()]
|
|
else:
|
|
names = range(len(arrays))
|
|
|
|
for array, section, name in zip(arrays, sections, names):
|
|
info = {}
|
|
|
|
info['name'] = name
|
|
info['shape'] = str(array.shape)
|
|
info['min'] = '{:.3e}'.format(section.min())
|
|
info['mean'] = '{:.3e}'.format(section.mean())
|
|
info['max'] = '{:.3e}'.format(section.max())
|
|
info['range'] = '{:.3e}'.format(section.max() - section.min())
|
|
info['height'] = section.shape[0]
|
|
|
|
infos.append(info)
|
|
|
|
write_pickle(infos, '{}/{}'.format(self.logdir, SECTION_INFO_FILENAME))
|
|
|
|
|
|
def build_frame(self, arrays):
|
|
self._maybe_clear_deque()
|
|
|
|
arrays = arrays if isinstance(arrays, list) else [arrays]
|
|
|
|
sections = self._arrays_to_sections(arrays)
|
|
self._save_section_info(arrays, sections)
|
|
final_image = self._sections_to_image(sections)
|
|
final_image = im_util.apply_colormap(final_image, self.config['colormap'])
|
|
|
|
return final_image
|
|
|
|
def update(self, config):
|
|
self.config = config
|