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應該也可以。