TensorBoard：将输出图像添加到回调

Question

我建立了一个网络，试图预测表面温度的光栅图像。 网络的输出是一个(1000, 1000) 大小的数组，代表一个光栅图像。对于训练和测试，这些将与各自样本的真实栅格进行比较。 我了解如何add the training image to my TensorBoard callback，但我还想将网络的输出图像添加到回调中，以便我可以直观地比较它们。这可能吗？

x = Input(shape = (2))
x = Dense(4)(x)
x = Reshape((2, 2))(x)

Reshape 将是最后一层（或某个反卷积层之前的一层）。

Answer 1

根据您使用的tensorflow 版本，我将提供 2 个不同的代码建议。我将假设您使用 > 2.0 并发布我用于该版本的图像到图像模型的代码。我基本上用嘈杂的图像初始化回调（我正在做去噪，但您可以轻松适应您的问题）和相应的地面真实图像。然后我使用模型在每个 epoch 之后进行推理。

"""Inspired by https://github.com/sicara/tf-explain/blob/master/tf_explain/callbacks/grad_cam.py"""
import numpy as np
import tensorflow as tf
from tensorflow.keras.callbacks import Callback


class TensorBoardImage(Callback):
    def __init__(self, log_dir, image, noisy_image):
        super().__init__()
        self.log_dir = log_dir
        self.image = image
        self.noisy_image = noisy_image

    def set_model(self, model):
        self.model = model
        self.writer = tf.summary.create_file_writer(self.log_dir, filename_suffix='images')

    def on_train_begin(self, _):
        self.write_image(self.image, 'Original Image', 0)

    def on_train_end(self, _):
        self.writer.close()

    def write_image(self, image, tag, epoch):
        image_to_write = np.copy(image)
        image_to_write -= image_to_write.min()
        image_to_write /= image_to_write.max()
        with self.writer.as_default():
            tf.summary.image(tag, image_to_write, step=epoch)

    def on_epoch_end(self, epoch, logs={}):
        denoised_image = self.model.predict_on_batch(self.noisy_image)
        self.write_image(denoised_image, 'Denoised Image', epoch)

所以通常你会这样使用它：

# define the model
model = Model(inputs, outputs)
# define the callback
image_tboard_cback = TensorBoardImage(
    log_dir=log_dir + '/images',
    image=val_gt[0:1],
    noisy_image=val_noisy[0:1],
)
# fit the model
model.fit(
    x, 
    y, 
    callbacks=[image_tboard_cback,],
)

如果您使用2.0 之前的版本，我可以直接转到我写的this gist（这有点复杂）。