SSD-pytorch训练自己的数据——狗

感谢：https://blog.csdn.net/After__today/article/details/108448263

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/1X5obnrBlcoTmjzPVlMwuOQ 
提取码：3cob 
复制这段内容后打开百度网盘手机App，操作更方便哦--来自百度网盘超级会员V7的分享