本章介绍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上留言即可。