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

作者：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的变量。

结语

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

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