SSD-pytorch训练自己的数据（附数据）

Windows下自己训练数据集，在pytorch上实现基于SSD的物体检测

（0）环境

Windows10

Anaconda3（python3.7.6）

GPU：GTX1660Ti-Max-Q

cuda10.0

cudnn7.6.5

pytorch1.2

（一）准备数据集

我用的labelImg标注，图片来源百度搜索图片：狗（共40张）。具体标注方法网上很多，不赘述。

数据集链接： 链接：https://pan.baidu.com/s/1We-XseyDVmSa-JWN1om0QQ 提取码：afqk

（二）下载源码并修改

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

1、data/init.py 第3行注释from .coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, COCO_ROOT, get_label_map

2、data/config.py 第5行HOME改为自己路径，第15行VOC num_classes = 类别数+1，第17行max_iter最大迭代次数

3、data/voc0712.py 第20行VOC_CLASSES改为自己的类，第29行VOC_ROOT如果VOCdevkit在data文件夹下就不用改，第93行改为image_sets=[('2007', 'trainval')]

4、layers/modules/multibox_loss.py 第97行的loss_c[pos] = 0前面加上一句loss_c = loss_c.view(num, -1)

5、ssd.py 所有的num_classes = 类别数+1(第32、198行)

6、train.py 第32行batch_size，第42行learning-rate根据自己电脑情况修改。搜索这里面的data[0]，全部替换为item()。第84、85行注释掉。第145行data_loader中的num_workers可以改为0避免出错。个人把151行iteration迭代次数的上限cfg['max_iter']加了1。第198行iteration % 5000 == 0，意味着每5000次保存一次模型，可改为500。

第195行创建txt记录loss值：

 with open('loss.txt', 'a') as f:

    
     f.write(str(loss.item()) + '\n')

第165行的images, targets = next(batch_iterator)改成：

 try:

    
     images, targets = next(batch_iterator)
    
 except StopIteration:
    
     batch_iterator = iter(data_loader)
    
     images, targets = next(batch_iterator)

（三）开始训练

开始训练时需要一个预训练文件vgg16_reducedfc.pth，百度云链接：https://pan.baidu.com/s/1pZGtpYYS1Q9TXwt59lhCFA
提取码：kn48

下载之后把他放在SSD项目下新建的weights文件夹下，然后就可以进行训练了。

数据集比较简单，我只训练了1000次。

训练了几十轮之后，loss降为nan，修改lr为0.0001，重新开始训练，loss值正常。

训练结束，编写loss_plot.py绘制loss变化情况：

（四）测试效果

新建test_image，在文件夹中放置几张待测图片，jupyter notebook中打开test.ipynb中测试效果：

很可惜，我手画的狗子它并不能识别出来。

完整代码、数据链接：https://pan.baidu.com/s/13mD5wto4T03ve6Mqktr3wg 提取码：bklg