【目标检测实战】Pytorch—SSD模型训练（VOC数据集

20.jpg

原文发表在：语雀文档

0.说明

SSD：Single Shot MultiBox Detector 和Yolo系列一样，他们都是one-stage系列的目标检测模型，SSD算法的官方实现是用的Caffe框架，源码在↓。有人将其改造成了Pytorch实现版，所以此代码可称为SSD算法的非官方Pytorch实现。

• SSD论文下载： 【2016】【SSD】1512.02325.pdf 软件翻译版：SSD.pdf
• caffe实现—官方：https://github.com/weiliu89/caffe/tree/ssd
• tensorflow实现—非官方：https://github.com/balancap/SSD-Tensorflow
• pytorch实现—非官方：https://github.com/amdegroot/ssd.pytorch

这里，需要说明：

• 1.目前只支持Python3.X版本
• 2.高度推荐通过支持CUDA的GPU来训练，用CPU也可以训练，不过非常非常慢
• 3.目前只支持VOC2007/VOC2012 、COCO数据集、对ImageNet数据集还不能够支持
• 4.本文的训练在windows10/ubuntu18.04都可正常跑通

推荐按照github首页给出的步骤一步步来：

image

整体来说，比较简单。不过，还是会遇到比较坑的地方，主要是visdom那里，这里直接跳过Visdom安装的步骤，因为即使安了，会发现代码也有报错，且visdom只是用于训练过程中可视化loss的页面，没有也不影响正常训练。

---

1.安装Pytorch

如果已安装Pytorch，可以直接跳过此步。由于我装了anaconda，所以直接用的conda安装，安装过程见pytorch官网

image

如果不翻墙，conda安装pytorch会比较慢，经常卡住，建议配置国内清华源,再下载速度会快很多

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

win10和ubuntu系统都建议直接用conda安装,步骤类似。安装完成在conda环境中输入python,进入python后如果import torch不报错则证明pytorch安装成功

2.资源准备

2.1 源码下载

git clone[https://github.com/amdegroot/ssd.pytorch.git](https://link.zhihu.com/?target=https%3A//github.com/amdegroot/ssd.pytorch.git)

或者下载zip包直接下载都可以ssd.pytorch-master.zip

2.2 VOC数据集下载

目前只支持COCO和VOC数据集，鉴于COCO数据集规模太大，这里仅准备VOC数据集。如果已下载好VOC数据集，可跳过此步。

VOC2007

训练集+测试集

http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar

VOC2012

训练集+测试集

http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar 将下载后的文件解压，将VOC2007和VOC2012文件夹共同放入VOCdevkit下，此文件夹路径即为后面代码中需要用到的VOC_ROOT。

例： VOC_ROOT='D:\personalProject\AI\pytorch\dataset\VOC\VOCdevkit' VOC_ROOT = '/home/lyon/AI/DeepLearning/ssd/dataset/VOCdevkit/'

2.3 模型文件下载

下载好的模型文件放入源码包weights文件夹下 vgg16_reducedfc.rar https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

3.安装visdom
visdom主要用于训练过程中，可通过浏览器可视化地观察loss情况和训练细节，如果不需要可视化地看训练过程，可跳过此步骤。
3.1安装visdom
pip install visdom
3.2修改源码
下载下来的visdom，如果直接用 python -m visdom.server 启动，浏览器访问http://localhost:8097/会发现页面什么都没有，因为在其源码server.py的1916行，会通过download_scripts()方法动态下载一些css、font、js等静态文件，如果不翻墙下载会失败导致文件缺失，进入页面后什么都没有....

image

解决办法： - 注释掉visdom包下server,py源码1917行的：download_scripts() - 手动将下载好的静态文件，放入visdom的static文件夹下
3.3增加静态文件
如果是conda安装的环境，可用conda env list查看各个环境的路径，然后找到visdom路径 visdom文件夹路径：conda环境路径/lib/site-packages/visdom

image

这里，可直接用我下载好的static.rar，替换掉原static文件，然后即可通过python -m visdom.server 正常启动，并访问http://localhost:8097/

4.模型训练

4.1设置模型参数

主要是在train.py中设置默认数据集，算法模型路径，batch-size，是否启动CUDA加速等等。这里源码默认的batch-size是32，根据自己GPU能力调节大小，我用GTX1060 6G的GPU发现32会崩，可以修改成16，RTX2060 SUPER下可以设置到32。详细配置如下：

# 加载模型初始参数
parser = argparse.ArgumentParser(
    description='Single Shot MultiBox Detector Training With Pytorch')
train_set = parser.add_mutually_exclusive_group()
# 默认加载VOC数据集
parser.add_argument('--dataset', default='VOC', choices=['VOC', 'COCO'],
                    type=str, help='VOC or COCO')
# 设置VOC数据集根路径
parser.add_argument('--dataset_root', default=VOC_ROOT,
                    help='Dataset root directory path')
# 设置预训练模型vgg16_reducedfc.pth
parser.add_argument('--basenet', default='vgg16_reducedfc.pth',
                    help='Pretrained base model')
# 设置批大小，根据自己显卡能力设置，默认为32，此处我改为16
parser.add_argument('--batch_size', default=16, type=int,
                    help='Batch size for training')
# 是否恢复中断的训练，默认不恢复
parser.add_argument('--resume', default=None, type=str,
                    help='Checkpoint state_dict file to resume training from')
# 恢复训练iter数，默认从第0次迭代开始
parser.add_argument('--start_iter', default=0, type=int,
                    help='Resume training at this iter')
# 数据加载线程数，根据自己CPU个数设置，默认为4
parser.add_argument('--num_workers', default=4, type=int,
                    help='Number of workers used in dataloading')
# 是否使用CUDA加速训练，默认开启，如果没有GPU，可改成False直接用CPU训练
parser.add_argument('--cuda', default=True, type=str2bool,
                    help='Use CUDA to train model')
# 学习率，默认0.001
parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,
                    help='initial learning rate')
# 最佳动量值，默认0.9（动量是梯度下降法中一种常用的加速技术，用于加速梯度下降，减少收敛耗时）
parser.add_argument('--momentum', default=0.9, type=float,
                    help='Momentum value for optim')
# 权重衰减，即正则化项前面的系数，用于防止过拟合；SGD，即随机梯度下降，多指mini-batch梯度下降
parser.add_argument('--weight_decay', default=5e-4, type=float,
                    help='Weight decay for SGD')
# gamma更新，默认值0.1
parser.add_argument('--gamma', default=0.1, type=float,
                    help='Gamma update for SGD')
# 使用visdom将训练过程loss图像可视化
parser.add_argument('--visdom', default=False, type=str2bool,
                    help='Use visdom for loss visualization')
# 权重保存位置，默认存在weights/下
parser.add_argument('--save_folder', default='weights/',
                    help='Directory for saving checkpoint models')
args = parser.parse_args()

设置好参数后，需要配置数据集加载路径，在源码包/data/voc0712.py中注释掉第28行，将VOC数据集路径替换为自己的路径：

image

4.2修改源码

由于Pytorch版本不同，较新版的代码直接运行会报错，需要修改部分代码，主要是将.data[0]的部分改成.item()

修改train.py

修改源码183.184两行

loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]

改为：

loc_loss += loss_l.item()
conf_loss += loss_c.item()

修改源码188行

print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')

改为：

print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')

修改源码165行

# load train data
images, targets = next(batch_iterator)

改为：

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

修改mutibox_loss.py

修改：源码包/layers/modules/mutibox_loss.py 调换第97,98行：

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

修改第114行为：

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

参考：
https://blog.csdn.net/qq_39571318/article/details/89788687 https://blog.csdn.net/dingkm666/article/details/88775428

修改coco.py

由于train.py会from data import ,而data初始化时会加载coco_labels.txt，这个文件在源码包中data/下，无论你是否下载了coco数据集都不影响其加载，加载时需要用到COCO_ROOT这个参数，需要修改COCO_ROOT为你的coco_labels.txt所在的父文件夹目录*。

4.3开始训练

训练方式以下两种都可以：

• 1.Pycharm中直接训练
• 2.cmd中进入conda的pytorch环境，python train.py

image

5.模型验证

5.1修改模型路径

第5步训练的模型，默认每迭代5000轮会在weights/下生成.pth模型文件，我这里训练了几个小时，迭代到25000轮。模型验证前，需要修改eval.py第39行，替换为你自己训练好的模型文件l路径。 生成的模型文件ssd300_COCO_25000.zip

image

5.2修改eval.py

运行eval.py，会对测试集中所有图片进行检测，并输出结果。这里实际运行发现windows系统下会有问题，需要做一点修改。 源码71行注释，改成如下形式：

#imgsetpath = os.path.join(args.voc_root, 'VOC2007', 'ImageSets','Main', '{:s}.txt')
imgsetpath = os.path.join(args.voc_root, 'VOC2007', 'ImageSets', 'Main') + os.sep + '{:s}.txt'

如不修改，则imgsetpath 不能正常拼接出来，导致报错

5.3运行eval.py验证模型

训练了几个小时，得到模型ssd300_COCO_25000.pth 验证结果如下：

Evaluating detections
Writing aeroplane VOC results file
Writing bicycle VOC results file
Writing bird VOC results file
Writing boat VOC results file
Writing bottle VOC results file
Writing bus VOC results file
Writing car VOC results file
Writing cat VOC results file
Writing chair VOC results file
Writing cow VOC results file
Writing diningtable VOC results file
Writing dog VOC results file
Writing horse VOC results file
Writing motorbike VOC results file
Writing person VOC results file
Writing pottedplant VOC results file
Writing sheep VOC results file
Writing sofa VOC results file
Writing train VOC results file
Writing tvmonitor VOC results file
VOC07 metric? Yes
Reading annotation for 1/4952
Reading annotation for 101/4952
Reading annotation for 201/4952
Reading annotation for 301/4952
......
......
Reading annotation for 4901/4952
Saving cached annotations to /home/lyon/AI/DeepLearning/ssd/dataset/VOCdevkit/VOC2007/annotations_cache/annots.pkl
AP for aeroplane = 0.7428
AP for bicycle = 0.8042
AP for bird = 0.6725
AP for boat = 0.6337
AP for bottle = 0.4359
AP for bus = 0.8039
AP for car = 0.8236
AP for cat = 0.8543
AP for chair = 0.5490
AP for cow = 0.7241
AP for diningtable = 0.6796
AP for dog = 0.8046
AP for horse = 0.8268
AP for motorbike = 0.7855
AP for person = 0.7249
AP for pottedplant = 0.4448
AP for sheep = 0.7060
AP for sofa = 0.7563
AP for train = 0.8028
AP for tvmonitor = 0.7311
Mean AP = 0.7153
~~~~~~~~
Results:
0.743
0.804
0.672
0.634
0.436
0.804
0.824
0.854
0.549
0.724
0.680
0.805
0.827
0.785
0.725
0.445
0.706
0.756
0.803
0.731
0.715
~~~~~~~~

--------------------------------------------------------------
Results computed with the **unofficial** Python eval code.
Results should be very close to the official MATLAB eval code.
--------------------------------------------------------------

Process finished with exit code 0

得到的mAP:0.7153