用SSD-Pytorch训练自己的数据集-完整教程

本机系统环境

• Ubuntu 18.04
• Miniconda3 4.6.14
• python 3.7.3
• PyTorch 1.0.1

SSD-PyTorch安装与测试

复制源码

GitHub: SSD: Single Shot MultiBox Object Detector, in PyTorch
git clone https://github.com/amdegroot/ssd.pytorch.git

下载数据集

COCO

sh data/scripts/COCO2014.sh

VOC 2007 & 2012

sh data/scripts/VOC2007.sh
sh data/scripts/VOC2012.sh
运行脚本后会把数据集存放在~/data中

修改源码

由于PyTorch版本差异，需要对源码几处地方进行修改：

• multibox_loss.py:

1. 调换第97,98行：

loss_c = loss_c.view(num, -1) 
loss_c[pos] = 0 # filter out pos boxes for now

2. 修改第114行为：

N = num_pos.data.sum().double() 
loss_l = loss_l.double() 
loss_c = loss_c.double()

下载权重文件

# 预训练权重
mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
# 训练好的VOC权重
wget https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth

测试

• 检测：
  用jupyter notebook运行demo/demo.ipynb
• 训练：
  python train.py
  默认训练数据集是VOC，需提前下载好数据集和预训练权重
• eval：
  python eval.py

训练自己的数据集

准备数据集

源码修改

训练与测试