class _TfWriter(_BaseWriter):
"""A class to write various TensorFlow data into TensorBoard summary files.
This class is intentionally not @traceable.
Args:
root_log_dir: The directory into which to store a new directory corresponding to this experiment's summary data
time_stamp: The timestamp of this experiment (used as a folder name within `root_log_dir`).
network: The network associated with the current experiment.
"""
tf_summary_writers: Dict[str, tf.summary.SummaryWriter]
def __init__(self, root_log_dir: str, time_stamp: str, network: TFNetwork) -> None:
super().__init__(root_log_dir=root_log_dir, time_stamp=time_stamp, network=network)
self.tf_summary_writers = DefaultKeyDict(
lambda key: (tf.summary.create_file_writer(os.path.join(root_log_dir, time_stamp, key))))
def write_epoch_models(self, mode: str) -> None:
with self.tf_summary_writers[mode].as_default(), summary_ops_v2.always_record_summaries():
summary_ops_v2.graph(backend.get_graph(), step=0)
for model in self.network.epoch_models:
summary_writable = (model.__class__.__name__ == 'Sequential'
or (hasattr(model, '_is_graph_network') and model._is_graph_network))
if summary_writable:
summary_ops_v2.keras_model(model.model_name, model, step=0)
def write_weights(self, mode: str, models: Iterable[Model], step: int, visualize: bool) -> None:
# Similar to TF implementation, but multiple models
with self.tf_summary_writers[mode].as_default(), summary_ops_v2.always_record_summaries():
for model in models:
for layer in model.layers:
for weight in layer.weights:
weight_name = weight.name.replace(':', '_')
weight_name = "{}_{}".format(model.model_name, weight_name)
with tfops.init_scope():
weight = backend.get_value(weight)
summary_ops_v2.histogram(weight_name, weight, step=step)
if visualize:
weight = self._weight_to_image(weight=weight, kernel_channels_last=True)
if weight is not None:
summary_ops_v2.image(weight_name, weight, step=step, max_images=weight.shape[0])
def close(self) -> None:
super().close()
modes = list(self.tf_summary_writers.keys()) # break connection with dictionary so can delete in iteration
for mode in modes:
self.tf_summary_writers[mode].close()
del self.tf_summary_writers[mode]
class _BaseWriter:
"""A class to write various types of data into TensorBoard summary files.
This class is intentionally not @traceable.
Args:
root_log_dir: The directory into which to store a new directory corresponding to this experiment's summary data
time_stamp: The timestamp of this experiment (used as a folder name within `root_log_dir`).
network: The network associated with the current experiment.
"""
summary_writers: Dict[str, SummaryWriter]
network: BaseNetwork
def __init__(self, root_log_dir: str, time_stamp: str,
network: BaseNetwork) -> None:
self.summary_writers = DefaultKeyDict(lambda key: (SummaryWriter(
log_dir=os.path.join(root_log_dir, time_stamp, key))))
self.network = network
def write_epoch_models(self, mode: str) -> None:
"""Write summary graphs for all of the models in the current epoch.
Args:
mode: The current mode of execution ('train', 'eval', 'test', 'infer').
"""
raise NotImplementedError
def write_weights(self, mode: str, models: Iterable[Model], step: int,
visualize: bool) -> None:
"""Write summaries of all of the weights of a given collection of `models`.
Args:
mode: The current mode of execution ('train', 'eval', 'test', 'infer').
models: A list of models compiled with fe.build whose weights should be recorded.
step: The current training step.
visualize: Whether to attempt to paint graphical representations of the weights in addition to the default
histogram summaries.
"""
raise NotImplementedError
def write_scalars(self, mode: str, scalars: Iterable[Tuple[str, Any]],
step: int) -> None:
"""Write summaries of scalars to TensorBoard.
Args:
mode: The current mode of execution ('train', 'eval', 'test', 'infer').
scalars: A collection of pairs like [("key", val), ("key2", val2), ...].
step: The current training step.
"""
for key, val in scalars:
self.summary_writers[mode].add_scalar(tag=key,
scalar_value=to_number(val),
global_step=step)
def write_images(self, mode: str, images: Iterable[Tuple[str, Any]],
step: int) -> None:
"""Write images to TensorBoard.
Args:
mode: The current mode of execution ('train', 'eval', 'test', 'infer').
images: A collection of pairs like [("key", image1), ("key2", image2), ...].
step: The current training step.
"""
for key, img in images:
if isinstance(img, ImgData):
img = img.paint_figure()
if isinstance(img, plt.Figure):
self.summary_writers[mode].add_figure(tag=key,
figure=img,
global_step=step)
else:
self.summary_writers[mode].add_images(
tag=key,
img_tensor=to_number(img),
global_step=step,
dataformats='NCHW'
if isinstance(img, torch.Tensor) else 'NHWC')
def write_embeddings(
self,
mode: str,
embeddings: Iterable[Tuple[str, Tensor, Optional[List[Any]],
Optional[Tensor]]],
step: int,
):
"""Write embeddings (like UMAP) to TensorBoard.
Args:
mode: The current mode of execution ('train', 'eval', 'test', 'infer').
embeddings: A collection of quadruplets like [("key", <features>, [<label1>, ...], <label_images>)].
Features are expected to be batched, and if labels and/or label images are provided they should have the
same batch dimension as the features.
step: The current training step.
"""
for key, features, labels, label_imgs in embeddings:
flat = to_number(reshape(features, [features.shape[0], -1]))
if not isinstance(label_imgs, (torch.Tensor, type(None))):
label_imgs = to_tensor(label_imgs, 'torch')
if len(label_imgs.shape) == 4:
label_imgs = permute(label_imgs, [0, 3, 1, 2])
self.summary_writers[mode].add_embedding(mat=flat,
metadata=labels,
label_img=label_imgs,
tag=key,
global_step=step)
def close(self) -> None:
"""A method to flush and close all connections to the files on disk.
"""
modes = list(self.summary_writers.keys(
)) # break connection with dictionary so can delete in iteration
for mode in modes:
self.summary_writers[mode].close()
del self.summary_writers[mode]
@staticmethod
def _weight_to_image(
weight: Tensor,
kernel_channels_last: bool = False) -> Optional[Tensor]:
"""Logs a weight as a TensorBoard image.
Implementation from TensorFlow codebase, would have invoked theirs directly but they didn't make it a static
method.
"""
w_img = squeeze(weight)
shape = backend.int_shape(w_img)
if len(shape) == 1: # Bias case
w_img = reshape(w_img, [1, shape[0], 1, 1])
elif len(shape) == 2: # Dense layer kernel case
if shape[0] > shape[1]:
w_img = permute(w_img, [0, 1])
shape = backend.int_shape(w_img)
w_img = reshape(w_img, [1, shape[0], shape[1], 1])
elif len(shape) == 3: # ConvNet case
if kernel_channels_last:
# Switch to channels_first to display every kernel as a separate images
w_img = permute(w_img, [2, 0, 1])
w_img = expand_dims(w_img, axis=-1)
elif len(shape) == 4: # Conv filter with multiple input channels
if kernel_channels_last:
# Switch to channels first to display kernels as separate images
w_img = permute(w_img, [3, 2, 0, 1])
w_img = reduce_sum(
abs(w_img),
axis=1) # Sum over the each channel within the kernel
w_img = expand_dims(w_img, axis=-1)
shape = backend.int_shape(w_img)
# Not possible to handle 3D convnets etc.
if len(shape) == 4 and shape[-1] in [1, 3, 4]:
return w_img