感谢: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的分享