如何在TensorBoard中记录图像
使用图像来直观的检视模型训练进展
封面图像：Photo by Ståle Grut on Unsplash
大型卷积神经网络的训练过程冗长且无趣，但是这是图像工程师工作的一部分。为了更好的监控训练进展，通常会将loss函数记录在日志文件中，并使用TensorBoard来查阅变化曲线。例如下图是BlazeFace的loss函数随epoch的变化曲线。
对于图像处理任务来说，除了具体的数值指标如loss外，模型在具体任务上的实际表现如何也可以作为训练过程的评估指标，而且更加直观易懂。这时候就需要在训练过程中添加模型的运行结果到训练日志中。TensorFlow提供了 tf.summary API来帮我们实现这一功能。具体的做法可以概括如下：
- 构建定制化的
callback函数。 - 在
callback函数中加载图像，馈入网络前向传播。 - 获取结果，绘制图像并存储。
以下内容以HRNet人脸特征点检测为例。如果你想看最终的代码，可以在GitHub上查阅。
yinguobing构建定制化的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栏目下找到当前模型的预测结果图像。
在你的项目中亲自试试吧！
Comments ()