SSD-Pytorch模型训练自己的数据集

开始之前声明:

博主用的环境是ubuntu18.04+anaconda3+pytorch1.4+pycharm

---

1.下载SSD-Pytorch代码

SSD-pytorch代码链接： https://github.com/amdegroot/ssd.pytorch

git clone https://github.com/amdegroot/ssd.pytorch

运行该代码下载到本地（如果下载太慢可以上传到码云，然后git clone码云地址）

2.准备数据集

没有数据集的同学可以下载代码自带的VOC和COCO数据集（./data/scripts目录下）

• 有自己的数据集请将数据集放置在./data目录下，例如VOC格式数据集，新建VOCdevkit文件夹，如下图所示，可以参考：https://blog.csdn.net/qq_34806812/article/details/81673798.
• 在Annotations中放置所有的标签，在JPEGimages中放置所有的图片，在ImagesSets/Main中放置train.txt/val.txt/test.txt（内容只有图片的名字，例如：00001,00002，不能带后缀jpg或者png）等，可以用脚本自己生成：https://blog.csdn.net/GeekYao/article/details/105074574.
  

3.根据自己的数据集修改代码

博主用的VOC格式的数据集，下面修改都是以VOC格式为例

3.1 config.py

• 找到config.py文件，打开修改VOC中的num_classes,根据自己的情况修改：classes+1（背景算一类），我这里就只有一类，所有是2
• 第一次调试最好修改一下max_iter,不然迭代次数太大，要好长时间
• 其他都是一些超参数，可以占时不修改

3.2VOC0712.py

• 根据自己的标签进行修改，博主这里只有一类，所以只有一个dargon fruit（注：如果只有一类，需要加上[ ]）
  
• image_sets中修改一下，根据自己的设置的数据集修改，我这里只有train和val

3.3 train.py

• 下载预训练模型。VGG16_reducedfc.pth
  链接：https://pan.baidu.com/s/1SsRN066zjmtWCbNClB6Z7Q 密码: jqs1
  
• 根据自己的显存修改batch_size，建议一开始修改小一点，博主1660ti 6G显存
  
• 将保存训练模型的参数调低一点，之前iter设置的1000，这里设置为500，之后根据自己情况在设置
• 顺便修改一下保存的模型名字，也可以之后修改，把COCO改成VOC，博主这里没修改

3.4 eval.py

• 添加训练好的模型到eval.py，对模型进行验证，我这里训练好的是ssd300_VOC_500.pth
• 将下面的

args = parser.parse_args()

修改为

args,unknow= parser.parse_known_args()

3.5 SSD.py

• 修改num_classes，跟上面config.py中的一致就行

---

修改完成后，运行train.py，完成训练之后，博主运行eval.py验证了训练的模型，AP只有63%，可能是博主数据集太少了

---

运行eval.py只能看到AP值，想要测试自己的图片，在jupyter notebook中运行demo.ipynb

将对应部分的代码，修改为以下这样即可，注意正确添加图片的路径

image = cv2.imread('../data/example3.jpg', cv2.IMREAD_COLOR