本章介紹SSD的Pytorch代碼實現。
具體原理參考上一章節:https://blog.csdn.net/IEEE_FELLOW/article/details/104241995
代碼我已經上傳我的github,歡迎star:https://github.com/midasklr/SSD.Pytorch
主要參考了https://github.com/amdegroot/ssd.pytorch代碼的實現,但是原作者代碼年久失修,並且僅僅實現了SSD300的結構。
在此基礎上,我修改支持最近幾個版本的Pytorch,並且參考caffe原實現實現了相應的SSD512。
我的環境如下:
- Ubuntu1604;
- python3.7 ;
- pytorch1.3;
- opencv等
數據集
我在PascalVOC數據集上訓練了SSD300和SSD512,並和原論文結果進行比較:
model | paper | this implements |
---|---|---|
SSD300 | 77.2 | 77.43 |
SSD512 | 79.8 | 79.80 |
模型文件:鏈接: https://pan.baidu.com/s/1DxlkOQzkFkkdYdNYsDx_MQ 提取碼: dd7m 複製這段內容後打開百度網盤手機App,操作更方便哦
當然你也可以訓練自己的數據集,這裏我使用SSD512訓練了一個車牌檢測,效果如下:
效果還行,就是延時還不理想,畢竟SSD結構還是比較老。但是這裏未做特徵融合還能識別出場景中小目標的車牌還是很驚人的。
訓練
具體見README。
測試
python eval.py --input 512 --trained_model weights/ssd512_VOC_73000_mAP79.80.pth
結果如下:
AP for aeroplane = 0.8861
AP for bicycle = 0.8694
AP for bird = 0.8078
AP for boat = 0.7698
AP for bottle = 0.6407
AP for bus = 0.8625
AP for car = 0.8825
AP for cat = 0.8671
AP for chair = 0.6424
AP for cow = 0.8712
AP for diningtable = 0.6781
AP for dog = 0.8572
AP for horse = 0.8781
AP for motorbike = 0.8531
AP for person = 0.8091
AP for pottedplant = 0.5479
AP for sheep = 0.8327
AP for sofa = 0.7562
AP for train = 0.8654
AP for tvmonitor = 0.7824
Mean AP = 0.7980
Demo
提供了demo.py,運行文件可以測試單張圖片:
結果
在VOC2007 測試集上:
model | paper | this implements |
---|---|---|
SSD300 | 77.2 | 77.43 |
SSD512 | 79.8 | 79.80 |
訓練過程:
訓練中MAP和epoch關係如下:
可以看到大約在250個循環左右達到最優秀,對比原作者Caffe中的參數發現基本吻合。
代碼已經上傳github:https://github.com/midasklr/SSD.Pytorch
歡迎大家star和訓練自己的數據集。有問題直接在issue上留言即可。