SSD（single-shot multibox detector）源码学习笔记

SSD（single-shot multibox detector）源码学习笔记

SSD是Wei Liu等人去年提出来的一个object detection框架，在PascalVOC上mAP可以超过著名的Faster RCNN，同时速度可以做到实时，简直强无敌。之前用别人的MXNet代码跑过实验，最近需要改进一下算法，所以去翻了一下原作者用Caffe实现的代码。

在看代码之前先简述一下SSD的做法：图片被送进网络之后先生成一系列feature map，传统一点的框架会在feature map（或者原图）上进行region proposal提取出可能有物体的部分然后进行分类，这一步可能非常费时，所以SSD就放弃了region proposal，而选择直接生成一系列defaul box（或者叫prior box），然后将这些default box回归到正确的位置上去，整个过程通过网络的一次前向传播就可以完成，非常方便，如下图所示。思路类似的算法还有YOLO，不过精度比SSD差了不少。

图1

原作者的代码分散在很多地方，包括 include/caffe/ 目录下面的 annotated_data_layer.hpp 、 detection_evaluate_layer.hpp 、 detection_output_layer.hpp 、 multibox_loss_layer.hpp 、 prior_box_layer.hpp ，以及对应的 src/caffe/layers/ 目录下面的cpp和cu文件，另外还有 src/caffe/utlis/ 目录下面的 bbox_util.cpp 。从名字就可以看出来， annotated_data_layer 是提供数据的、 detection_evaluate_layer 是验证模型效果用的、 detection_output_layer 是输出检测结果用的、 multibox_loss_layer 是loss、 prior_box_layer 是计算prior bbox的。其中跟训练有关的是 annotated_data_layer 、 multibox_loss_layer 、 prior_box_layer ，我就重点看这三个。

annotated_data_layer

这部分代码没有太多好说的，就是把图片读出来（包含一些data augmentation），同时把每张图里的groundtruth bbox读出来传给下一层，没做什么特殊处理。

prior_box_layer

这一层完成的是给定一系列feature map后如何在上面生成prior box。SSD的做法很有意思，对于输入大小是 W×H 的feature map，生成的prior box中心就是 W×H 个，均匀分布在整张图上，像下图中演示的一样。在每个中心上，可以生成多个不同长宽比的prior box，如[1/3, 1/2, 1, 2, 3]。所以在一个feature map上可以生成的prior box总数是 W×H×length_of_aspect_ratio ，对于比较大的feature map，如VGG的conv4_3，生成的prior box可以达到数千个。当然对于边界上的box，还要做一些处理保证其不超出图片范围，这都是细节了。

这里需要注意的是，虽然prior box的位置是在 W×H 的格子上，但prior box的大小并不是跟格子一样大，而是人工指定的，原论文中随着feature map从底层到高层，prior box的大小在0.2到0.9之间均匀变化。

图2

一开始看SSD的时候很困扰我的一点就是形状的匹配问题：SSD用卷积层做bbox的拟合，输出的不应该是feature map吗，怎么能正好输出4个坐标呢？这里的做法有点暴力，比如需要输出 W×H×length_of_aspect_ratio×4 个坐标，就直接用 length_of_aspect_ratio×4 个channel的卷积层做拟合，这样就得到 length_of_aspect_ratio×4 个大小为 W×H 的feature map，然后把feature map拉成一个长度为 W×H×length_of_aspect_ratio×4 的向量，用SmoothL1之类的loss去拟合，效果还意外地不错……

multibox_loss_layer

FindMatches

我们已经在图上画出了prior box，同时也有了ground truth，那么下一步就是将prior box匹配到ground truth上，这是在 src/caffe/utlis/bbox_util.cpp 的 FindMatches 函数里完成的。值得注意的是这里不光是给每个groudtruth box找到了最匹配的prior box，而是给每个prior box都找到了匹配的groundtruth box（如果有的话），这样显然大大增大了正样本的数量。

MineHardExamples

给每个prior box找到匹配（包括物体和背景）之后，似乎可以定义一个损失函数，给每个prior box标记一个label，扔进去一通训练。但需要注意的是，任意一张图里负样本一定是比正样本多得多的，这种严重不平衡的数据会严重影响模型的性能，所以对prior box要有所选择。这一步是在 src/caffe/utlis/bbox_util.cpp 的 MineHardExamples 函数里完成的，我没细看= =。

EncodeLocPrediction && EncodeConfPrediction

因为我们对prior box是有选择的，所以数据的形状在这里已经被打乱了，没办法直接在后面连接一个loss（Caffe等框架需要每一层的输入是四维张量），所以需要我们把选出来的数据重新整理一下，这一步是在 src/caffe/utlis/bbox_util.cpp 的 EncodeLocPrediction 和 EncodeConfPrediction 两个函数里完成的。

训练

Wei Liu等人在这里的实现方法是在multibox_loss_layer里又实例化了两个loss_layer，前向传播的时候需要先经过前面三步，然后把数据手动喂给loss，反向传播的时候也需要手动把残差从loss取出来写回给前面的卷积层。

思考

我C++水平很一般，看上面的Caffe代码看得很痛苦，看代码的时候一直在想有没有更优雅的实现方法，哪些事情一定要自己做，哪些事情可以利用框架里现成的模块，Python和C++之间如何权衡，诸如此类的问题。

首先，我觉得上面提到的三步都不一定要用C++完成，用Python实现得好一点效率应该不会差太多，毕竟训练NN大头还是在卷积等等计算上（没错我喜欢用Python）。其次，我觉得可以不用把两个loss写在multibox_loss里，输出两个blob给下一层就好了，这样可以让一个模块就做一件事，模型结构看起来也更清晰一点。另外，从这篇文章来看，貌似TensorFlow可以只用Python实现SSD（negative mining被简化了），可读性更高一点，可惜我不懂TensorFlow。我猜MXNet应该也可以。