如何用神经网络“寻找威利”

作者:Tadej Magajna

编译:Bing

《威利在哪里?》(Where’s Wally)是由英国插画家马丁·汉德福特(Martin Handford)创作的一套儿童绘本。这个书的目标就是在一张人山人海的图片中找出一个特定的人物——威利(Wally)。“Where’s Wally”的商标已在28个国家进行了注册,为方便语言翻译,每一个国家都会给威利起一个新名字,最成功的是北美版的“Where’s Waldo”,在这里,威利改名成了沃尔多(Waldo)。

现在,机器学习博主Tadej Magajna另辟蹊径,利用深度学习解开“威利在哪里”的问题。与传统的计算机视觉图像处理方法不同的是,它只使用了少数几个标记出威利位置的图片样本,就训练成了一套“寻找威利”的系统。

训练过的图像评估模型和检测脚本发布在作者的GitHub repo上。

本文介绍了用TensorFlow物体检测API训练神经网络、并用相应的Python脚本寻找威利的过程。大致分为以下几步:

将图片打标签后创建数据集,其中标签注明了威利在图片中的位置,用x,y表示;

用TensorFlow物体检测API获取并配置神经网络模型;

在数据集上训练模型;

用导出的图像测试模型;

开始前,请确保你已经按照说明安装了TensorFlow物体检测API。

创建数据集

虽说深度学习中最重要的环节是处理神经网络,但不幸的是,数据科学家们总要花费大量时间准备训练数据。

最简单的机器学习问题最终得到的通常是一个标量(如数字检测器)或是一个分类字符串。TensorFlow物体检测API在训练数据是则将上述两个结果结合了起来。它由一系列图像组成,并包含目标对象的标签和他们在图像中的位置。由于在二维图像中,两个点足以在对象周围绘制边界框,所以图像的定位只有两个点。

为了创建训练集,我们需要准备一组Where’s Wally的插画,并标出威利的位置。在此之前已经有人做出了一套解出威利在哪里的训练集。

最右边的四列描述了威利所在的位置

创建数据集的最后一步就是将标签(.csv)和图片(.jpeg)打包,存入单一二分类文件中(.tfrecord)。详细过程可参考这里,训练和评估过程也可以在作者的GitHub上找到。

准备模型

TensorFlow物体检测API提供了一组性能不同的模型,它们要么精度高,但速度慢,要么速度快,但精度低。这些模型都在公开数据集上经过了预训练。

虽然模型可以从头开始训练,随机初始化网络权重,但这可能需要几周的时间。相反,这里作者采用了一种称为迁移学习(Transfer Learning)的方法。

这种方法是指,用一个经常训练的模型解决一般性问题,然后再将它重新训练,用于解决我们的问题。也就是说,与其从头开始训练新模型,不如从预先训练过的模型中获取知识,将其转移到新模型的训练中,这是一种非常节省时间的方法。

作者使用了在COCO数据集上训练过的搭载Inception v2模型的RCNN。该模型包含一个.ckpycheckpoint文件,可以利用它开始训练。

配置文件下载完成后,请确保将“PATHTOBE_CONFIGURED”字段替换成指向checkpoint文件、训练和评估的.tfrecord文件和标签映射文件的路径。

最后需要配置的文件是labels.txt映射文件,其中包含我们所有不同对象的标签。由于我们寻找的都是同一个类型的对象(威利),所以标签文件如下:

最终应该得到:

一个有着checkpoint文件的预训练模型;

经过训练并评估的.tfrecord数据集;

标签映射文件;

指向上述文件的配置文件。

然后就可以开始训练啦。

训练

TensorFlow物体检测API提供了一个十分容易上手的Python脚本,可以在本地训练模型。它位于models/research/object_detection中,可以通过以下命令运行:

PATH_TO_PIPELINE_CONFIG是通往配置文件的路径,PATH_TO_TRAIN_DIR是新创建的directory,用来储存checkpoint和模型。

train.py的输出看起来是这样:

用最重要的信息查看是否有损失,这是各个样本在训练或验证时出现错误的总和。当然,你肯定希望它降得越低越好,因为如果它在缓慢地下降,就意味着你的模型正在学习(要么就是过拟合了你的数据……)。

你还可以用Tensorboard显示更详细的训练数据。

脚本将在一定时间后自动存储checkpoint文件,万一计算机半路崩溃,你还可以恢复这些文件。也就是说,当你想完成模型的训练时,随时都可以终止脚本。

但是什么时候停止学习呢?一般是当我们的评估集损失停止减少或达到非常低的时候(在这个例子中低于0.01)。

测试

现在,我们可以将模型用于实际测试啦。

首先,我们需要从储存的checkpoint中输出一个推理图(interference graph),利用的脚本如下:

产生的推理图就是用来Python脚本用来找到威利的工具。

作者写了几个简单目标定位的脚本,其中find_wally.py和find_wally_pretty.py都可以在他的GitHub上找到,并且运行起来也很简单:

或者

不过当你在自己的模型或图像上运行脚本时,记得改变model-path和image-path的变量。

结语

模型的表现出乎意料地好。它不仅从数据集中成功地找到了威利,还能在随机从网上找的图片中找到威利。

但是如果威利在图中特别大,模型就找不到了。我们总觉得,不应该是目标物体越大越好找吗?这样的结果表明,作者用于训练的图像并不多,模型可能对训练数据过度拟合了。

  • 发表于:
  • 原文链接:http://kuaibao.qq.com/s/20171208G113BS00?refer=cp_1026

同媒体快讯

相关快讯

扫码关注云+社区