SSD模型训练技巧小结

最新推荐文章于 2024-08-08 08:20:50 发布

Jesse_Mx

最新推荐文章于 2024-08-08 08:20:50 发布

阅读量8.5k

点赞数 6

分类专栏： SSD 文章标签： ssd 训练

本文链接：https://blog.csdn.net/jesse_mx/article/details/78928262

版权

前言

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'