Press "Enter" to skip to content

如何在TensorBoard中记录图像

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

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

 

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

loss函数随epoch的变化曲线

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

 

callback
callback

 

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

 

yinguobing/facial-landmark-detection-hrnet

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

yinguobing GitHub

构建定制化的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初始化代码

目前这个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栏目下找到当前模型的预测结果图像。

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

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注