PyTorch 1.0 正式版刚出，这里就有人放出了 SSD 的高效的实现方案。亮点如下：
- PyTorch 1.0
- GPU/CPU NMS
- 多 GPU 训练
- 模块化
- 可视化(支持 Tensorboard)
- 支持 CPU 推理
文章目录
安装
系统要求
- Python3
- PyTorch 1.0
- yacs
- GCC >= 4.9
- OpenCV
编译
1 2 3 | # build nms cd ext python build.py build_ext develop |
训练
准备数据库
Pascal VOC
需要使用如下的目录结构：
1 2 3 4 5 6 7 8 9 10 11 12 | VOC_ROOT |__ VOC2007 |_ JPEGImages |_ Annotations |_ ImageSets |_ SegmentationClass |__ VOC2012 |_ JPEGImages |_ Annotations |_ ImageSets |_ SegmentationClass |__ ... |
COCO
需要使用如下的目录结构：
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | COCO_ROOT |__ annotations |_ instances_valminusminival2014.json |_ instances_minival2014.json |_ instances_train2014.json |_ instances_val2014.json |_ ... |__ train2014 |_ <im-1-name>.jpg |_ ... |_ <im-N-name>.jpg |__ val2014 |_ <im-1-name>.jpg |_ ... |_ <im-N-name>.jpg |__ ... |
单 GPU 训练
1 2 | # for example, train SSD300: python train_ssd.py --config-file configs/ssd300_voc0712.yaml --vgg vgg16_reducedfc.pth |
多 GPU 训练
1 2 3 | # for example, train SSD300 with 4 GPUs: export NGPUS=4 python -m torch.distributed.launch --nproc_per_node=$NGPUS train_ssd.py --config-file configs/ssd300_voc0712.yaml --vgg vgg16_reducedfc.pth |
预训练模型
| Weights | |
|---|---|
| SSD300* | ssd300_voc0712_mAP77.83.pth(100 MB) |
| SSD512* | ssd512_voc0712_mAP80.25.pth(104 MB) |
性能
论文中的数据:
| VOC2007 test | |
|---|---|
| SSD300* | 77.2 |
| SSD512* | 79.8 |
本实现的结果:
| VOC2007 test | |
|---|---|
| SSD300* | 77.8 |
| SSD512* | 80.2 |
Github 地址
https://github.com/lufficc/SSD
本站微信群、QQ群(三群号 726282629)：
