如何在TensorBoard中记录图像

封面图像：Photo by Ståle Grut on Unsplash

大型卷积神经网络的训练过程冗长且无趣，但是这是图像工程师工作的一部分。为了更好的监控训练进展，通常会将loss函数记录在日志文件中，并使用TensorBoard来查阅变化曲线。例如下图是BlazeFace的loss函数随epoch的变化曲线。

对于图像处理任务来说，除了具体的数值指标如loss外，模型在具体任务上的实际表现如何也可以作为训练过程的评估指标，而且更加直观易懂。这时候就需要在训练过程中添加模型的运行结果到训练日志中。TensorFlow提供了 tf.summary API来帮我们实现这一功能。具体的做法可以概括如下：

1. 构建定制化的 callback 函数。
2. 在 callback 函数中加载图像，馈入网络前向传播。
3. 获取结果，绘制图像并存储。

以下内容以HRNet人脸特征点检测为例。如果你想看最终的代码，可以在GitHub上查阅。

yinguobing/facial-landmark-detection-hrnet

A TensorFlow implementation of HRNet for facial landmark detection. - yinguobing/facial-landmark-detection-hrnet

GitHubyinguobing

构建定制化的callback类

TensorFlow使用callback类来实现训练过程中的人工干预。对于我们即将实现的callback对象，我们需要提供用于预测的样本图像以及日志路径。同时实例化一个日志写入器方便后续使用。

class LogImages(keras.callbacks.Callback):
    def __init__(self, logdir, sample_image):
        super().__init__()
        self.file_writer = tf.summary.create_file_writer(logdir)
        self.sample_image = sample_image

目前这个callback对象没有任何实际功能。我们需要根据需求来定义特定的方法。

实现图像日志方法

考虑到用户与训练过程的交互可能发生的时间节点，TensorFlow提供了大量预设方法。假设我们需要在每个epoch结束后绘制当前模型的预测结果，则需要利用callback类的 on_epoch_end 函数。该函数会在每个epoch结束后被调用。

def on_epoch_end(self, epoch, logs={}):

在函数中，分步实现以下几个功能。

读取图像并预处理

你可以使用任何你喜欢的方式，这里我使用OpenCV实现图像预处理。图像读入后转换为RGB格式，尺寸变化为网络支持的大小，并按照模型定义实施归一化处理。我保留了读入的原始图像以便后续写入到日志中。

# Read in the image file.
image = cv2.imread(self.sample_image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
img = cv2.resize(image, (256, 256))
img = normalize(img)

馈入网络前向传播

一旦样本准备完成，便可以通过 self.model 获得当前正在训练的模型。这是一个Keras模型，所以可以使用模型的 predict 方法得到预测结果。

heatmaps = self.model.predict(tf.expand_dims(img, 0))[0]

解析结果并记录

获取模型预测结果后，便可以根据自己的需求将结果绘制在原始图像中。通常情况下解析过程会被打包成函数以供调用。

marks, _ = parse_heatmaps(heatmaps, image.shape[:2])
	for mark in marks:
    	cv2.circle(image, tuple(mark.astype(int)), 2, (0, 255, 0), -1)

此刻，面部特征点的检测结果已经绘制在了image上。接下来保存该图像到训练日志。

with self.file_writer.as_default():
    img_tensor = tf.expand_dims(image, 0)
    tf.summary.image("test-sample", img_tensor, step=epoch)

然后开始模型训练，你会在TensorBoard中的IMAGE栏目下找到当前模型的预测结果图像。

在你的项目中亲自试试吧！