深度學習目標檢測系列論文閱讀和Pytorch實現(三)——SSD代碼實現

本章介紹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。

我的環境如下:

  1. Ubuntu1604;
  2. python3.7 ;
  3. pytorch1.3;
  4. 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上留言即可。

發佈了12 篇原創文章 · 獲贊 6 · 訪問量 3萬+
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章