SSD模型训练技巧小结

前言

SSD是目前先进的one-step目标检测算法，针对该框架的训练还有不少技巧可以挖掘，本文试着写写一些常见和不常见的技巧，在不改变网络架构的条件下（相反则是DSSD，R-SSD等改动卷积层结构的系列算法），应该能对精度提升有所帮助。这里只是抛砖引玉，不敢说真的很懂。

抽取权重

一般而言，我们都是用预训练模型来开始新的数据集训练，常用的预训练模型主要是Imagenet，Pascal VOC，COCO上训练好的，这些都能在网上下载到。

那现在如果我们训练KITTI数据集，它只有4类（car，person，cyclist，bg），如果使用预训练模型直接开始训练，那么预训练模型（比如VOC和COCO）中所有mbox_conf层的参数都被浪费了，必须重新学习，这是因为类别数量不一致，卷积核参数无法继承。这里mbox_loc层还会保留，因为SSD算法共享位置。

考虑到VOC数据集和COCO数据集中都含有汽车、行人和自行车这三类，因此可以考虑从这两类预训练模型中抽取这4类（加上背景）的权重，形成一个仅针对这几类目标的caffemodel模型。如此这般，排除了其他无关类型的参数，这个新模型的参数就能得到充分利用，再基于该模型进行finetune，其效果也会更好一些。

博主参考pycaffe相关接口写了一个简单的抽取程序，这个针对的是VOC训练出来的模型，如果想针对COCO模型，则稍加修改就行。写的比较仓促，代码效率未必最高，敬请见谅。

# convert_ssd_caffemodel.py
# coding:utf-8
import numpy as np
import sys

caffe_root='/home/mx/caffe/'
sys.path.insert(0,caffe_root+'python')
import caffe

caffe.set_device(0)
caffe.set_mode_gpu() # 使用GPU模式

voc_deploy='/home/mx/tempfile/net_surgery/deploy_voc.prototxt' # 21类的deploy文件
voc_caffemodel='/home/mx/tempfile/net_surgery/VGG_VOC0712_SSD_300x300_iter_120000.caffemodel' # 21类的caffemodel文件
kitti_deploy='/home/mx/tempfile/net_surgery/deploy_kitti.prototxt' # 4类的deploy文件

voc_net=caffe.Net(voc_deploy,voc_caffemodel,caffe.TEST) # 从caffemodel加载参数
kitti_net=caffe.Net(kitti_deploy,caffe.TEST) # 参数由系统初始化

# 需要修改的层，从21类中挑出4类使用
change_layer=['conv4_3_norm_mbox_conf','fc7_mbox_conf','conv6_2_mbox_conf','conv7_2_mbox_conf','conv8_2_mbox_conf','conv9_2_mbox_conf']

# 定义字典，层名和输入通道数对应起来，有两种类型，prior_box=4 or 6
name_channel_84={'conv4_3_norm_mbox_conf':512,'conv8_2_mbox_conf':256,'conv9_2_mbox_conf':256}
name_channel_126={'fc7_mbox_conf':1024,'conv6_2_mbox_conf':512,'conv7_2_mbox_conf':256}

for name in voc_net.params.keys():
if(name not in change_layer):
if(name=='conv4_3_norm'):
kitti_net.params[name][0].data[...] = voc_net.params[name][0].data # 单独处理conv4_3_norm
else:
kitti_net.params[name][0].data[...] = voc_net.params[name][0].data
kitti_net.params[name][1].data[...] = voc_net.params[name][1].data

# 处理mbox_conf的卷积层参数，首先是prior_box=4的三个层
for name,channel in name_channel_84.items():
weight=voc_net.params[name][0].data # 暂存卷积核权重
bias=voc_net.params[name][1].data # 暂存偏置项
w=np.zeros((16,channel,3,3)) # 目标维度是(16,channel,3,3)，初始化为全0
b=np.zeros((16,)) # 偏置项是一维矩阵
# 按照voc的类别顺序，依次提取背景(0)、自行车(2)、汽车(7)、行人(15)的卷积核权重和偏置项，排布方式要清楚。
for i in [0,1,2,3]:
w[4 * i, ...] = weight[21 * i, ...]
w[1 + 4 * i, ...] = weight[2 + 21 * i, ...]
w[2 + 4 * i, ...] = weight[7 + 21 * i, ...]
w[3 + 4 * i, ...] = weight[15 + 21 * i, ...]
b[4 * i] = bias[21 * i]
b[1 + 4 * i] = bias[2 + 21 * i]
b[2 + 4 * i] = bias[7 + 21 * i]
b[3 + 4 * i] = bias[15 + 21 * i]
# 存入到kitti_net的对应层
kitti_net.params[name][0].data[...]=w
kitti_net.params[name][1].data[...]=b

# 处理mbox_conf的卷积层参数，首然后是prior_box=6的三个层
for name,channel in name_channel_126.items():
weight=voc_net.params[name][0].data
bias=voc_net.params[name][1].data
w=np.zeros((24,channel,3,3))
b=np.zeros((24,))
for i in [0,1,2,3,4,5]:
w[4 * i, ...] = weight[21 * i, ...]
w[1 + 4 * i, ...] = weight[2 + 21 * i, ...]
w[2 + 4 * i, ...] = weight[7 + 21 * i, ...]
w[3 + 4 * i, ...] = weight[15 + 21 * i, ...]
b[4 * i] = bias[21 * i]
b[1 + 4 * i] = bias[2 + 21 * i]
b[2 + 4 * i] = bias[7 + 21 * i]
b[3 + 4 * i] = bias[15 + 21 * i]
kitti_net.params[name][0].data[...]=w
kitti_net.params[name][1].data[...]=b

print('all the parameters have been wrote.')

# 保存caffemodel文件到本地目录
# 使用该模型作为预训练模型，可加快训练速度，精度也有一定提升
kitti_net.save('/home/mx/tempfile/net_surgery/VGG_VOC0712_SSD_300x300_iter_120000_4class.caffemodel')
print('the caffemodel has been saved!')

程序及其所需文件也可到此处下载：convert_ssd_caffemodel ，至于

VGG_VOC0712_SSD_300x300_iter_120000.caffemodel

，请到SSD项目主页下载。

原caffemodel的大小为105.21m，转换后大小为96.1MB，使用抽取后的caffemodel进行图片检测，效果尚可。



然后使用

VGG_VOC0712_SSD_300x300_iter_120000_4class.caffemodel

作为预训练模型来训练KITTI数据集，batch_size=32，发现达到同样的mAP，差不多能提前5000~10000次迭代，说明此模型能有效加快拟合。至于mAP，在没有其他改动的情况下，增加了大约1.5%，如果再仔细调节学习率，可能会达到2%~4%的提升。用COCO的转换模型，效果则还要更好一些，毕竟COCO数据集规模更大。