基于自制数据集的MobileNet-SSD模型训练

关注

转载最新推荐文章于 2026-05-23 11:46:37 发布 1.4k 阅读 AI 写同款· GEO 优化 ›

SSD与MobileNet详解 1.概述本文档阐述SSD检测算法原理，及以MobileNet为Backbone的网络搭建方式。 Why SSD？无论是机器学习或是深度学习一般都可分为两个部分：特征提取与分类任务。在传统的机器学习方法中，特征提取需要依据图像以及特有的检测目的抓取特有特征，如偏重物体轮廓的HOG特征，注重明暗对比的Haar特征等，特征被描述之后送入机器学习算法分类，如SVM、Adaboost等，进而判断物... 阅读详情

基于自制数据集的MobileNet-SSD模型训练

来源：QQ快报责任编辑:小易

“本文主要内容：基于自制的仿VOC数据集，利用caffe框架下的MobileNet-SSD模型训练。”

本文的base是https://github.com/chuanqi305/MobileNet-SSD，这个project又是基于https://github.com/weiliu89/caffe/tree/ssd，因此项目编译和数据库生成大多同 weiliu89的base。以下从环境搭建、数据集制作、模型训练、模型测试四个环节介绍整个过程。

—

环境搭建

weiliu89的caffe框架下SSD环境搭建过程参见https://github.com/weiliu89/caffe/tree/ssd和网络相关资料，这里不再赘述，主要流程包括下载代码、修改配置文件、编译源代码等。编译通过之后就可以玩模型啦。

—

数据集制作

网络上大多数资料都是在介绍如何利用VOC2007和VOC2012数据集开始训练，本文介绍的是制作自己的仿VOC数据集，对接工程实际。

准备数据集

在caffe根目录/data下新建一个文件夹（下文以VOCdevkit2文件夹名为例），该文件夹下新建子文件夹Annotations、ImageSets和JPEGImages。JPEGImages文件夹存放图片，Annotations文件夹下存放与JPEGImages中图片一一对应的xml文件，ImageSets文件夹下包含Main子文件夹，Main里面存有train.txt和test.txt两个文本文件，这两个文件里面的每一行分别是预作为训练集和测试集的文件名。以上的文件结构参考了VOCdevkit，但是去除了2007和2012两个子文件夹的设计，简化了后续数据处理的代码编写。

xml标注文件通过labelImg工具获取，前提是制定好将要分类的类别，以及投入人力标注图片。

制作map.prototxt

在caffe根目录/data/VOC0712下创建一个自己的labelmap.prototxt文件，该文件中的类别为“准备数据集”部分的类别数+1（background）。笔者的数据集共9类，加上background共计10类，下图展示了前5类。

创建lmdb数据库

参考weiliu89的caffe框架下SSD项目，可将caffe根目录/data/VOC0712下原先的create_list.sh和create_data.sh两个脚本复制一份，修改名称后修改脚本文件内容。create_list.sh需要修改root_dir为自己的路径，同时去掉原先2007和2012两个子文件夹的处理逻辑；create_data.sh需要修改data_root_dir、mapfile为自己的路径。

修改完成后，在caffe根目录下运行：

./data/VOC0712/create_list.sh

./data/VOC0712/create_data.sh

创建完成后，应该在caffe根目录/data/VOCdevkit2/VOC0712/lmdb下产生VOC0712_test_lmdb和VOC0712_train_lmdb两个文件夹，内置生成的数据库文件。

—

模型训练

直到“模型训练”环节，我们才开始接触 chuanqi305的MobileNet-SSD里的code。按照 chuanqi305的readme，将MobileNet-SSD文件夹放入caffe根目录/examples，并使用fanqiang技术下载预训练模型，做好模型训练的准备工作。

生成prototxt文件

weiliu89的caffe框架下SSD是利用python脚本ssd_pascal.py自动生成prototxt文件并开始训练的，而chuanqi305的MobileNet-SSD则是利用gen_model.sh脚本生成prototxt文件,使用train.sh脚本开始训练。

在caffe根目录/examples/MobileNet-SSD下运行

./gen_model.sh NUM（NUM为你自己的类别数，同labelmap.prototxt，笔者这里NUM=10）

即可在caffe根目录/examples/MobileNet-SSD/example下生成MobileNetSSD_train.prototxt、MobileNetSSD_test.prototxt和MobileNetSSD_deploy.prototxt。

三个prototxt文件生成之后，需要做如下修改：

MobileNetSSD_train.prototxt、MobileNetSSD_test.prototxt中，将data_param层里的source字段改写成项目实际lmdb的路径（按理说加了软链接后的路径不必修改，但保险起见这里可以修改为绝对路径），同时可修改batch_size的大小（可使用默认值）。

检查所有prototxt文件的num_classes是否为你的分类数。num_classes是脚本读取labelmap.prototxt生成的，这里可以查看一下是否被修改正确。

prototxt很长，可以使用caffe根目录/python/drawnet.py画出模型，这里先挖个坑，争取看懂算法模型后再叙。

训练模型

接下来进入正题，在caffe根目录/examples/MobileNet-SSD下运行

./train.sh

开启训练的大幕。./train.sh里面的内容比较简单：

第2行是train.prototxt的路径，第7行是snapshot保存中间模型的路径，第8行是slover文件的路径，第9行是预训练权重，第10行是用到的gpu编号，这些都可以按需修改。笔者将最后一行的gpu从0修改为0,1,可利用双GPU训练。

solver_train.prototxt文件内容大多同SSD中的solver.prototxt，读者可根据需要调整：

train_net:”example/MobileNetSSD_train.prototxt” #训练用的网络文件

test_net:”example/MobileNetSSD_test.prototxt” #测试用的网络文件

test_iter: 673 #测试时的迭代次数

test_interval: 1000 #每训练1000次迭代一次

base_lr: 0.0005 #基准学习率

display: 5 #每5次迭代屏幕显示一次信息

max_iter: 20000 #最大迭代次数

lr_policy: “multistep” #学习率调整策略

gamma: 0.5

weight_decay: 0.00005 #权重衰减

snapshot: 1000 #每1000次迭代抓取一次快照

snapshot_prefix:”snapshot_power/mobilenet” #快照保存的前缀

solver_mode: GPU #这里一定要GPU

debug_info: false

snapshot_after_train: true

test_initialization: false #不做测试初始化

average_loss: 10 #每10次迭代平均下loss

stepvalue: 10000

stepvalue: 20000

iter_size: 1

type: “RMSProp” #一种优化方法，挖坑，以后填

eval_type: “detection” #这里一定要写检测，默认为classification

ap_version: “11point” #计算AP的方法

如果运气好的话，运行./train.sh脚本后就可以愉快地训练了，不过假如和笔者一样点背的话，可能在训练中遇到如下的错误：

OpenCV Error: Assertion failed ((scn == 3 scn == 4) && (depth == CV_8U depth == CV_32F)) in ipp_cvtColor,file /home/user1/opencv-3.1.0/modules/imgproc/src/color.cpp, line 7646

terminate called after throwing an instanceof ‘cv::Exception’

what(): /home/user1/opencv-3.1.0/modules/imgproc/src/color.cpp:7646:error: (-215) (scn == 3 scn == 4) && (depth == CV_8U depth ==CV_32F) in function ipp_cvtColor

或者

16:07:24.865304 12717data_transformer.cpp:621] Check failed: mean_values_.size() == 1 mean_values_.size() == img_channels Specify either 1 mean_value or as many aschannels: 1

或者

18:12:50.348572 40896annotated_data_layer.cpp:205] Check failed: std::equal(top_shape.begin() + 1,top_shape.begin() + 4, shape.begin() + 1)

解决方案：

MobileNetSSD_train.prototxt和MobileNetSSD_test.prototxt中找到transform_param层，添加force_color:true。

参考资料http://blog.youkuaiyun.com/u013250416/article/details/78676784?locationNum=7&fps=1

解决上述问题后，就可以愉快地训练了。贴一张训练初期的截图。