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

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

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

01

—

环境搭建

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

02

—

数据集制作

网络上大多数资料都是在介绍如何利用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两个文件夹，内置生成的数据库文件。

03

—

模型训练

直到“模型训练”环节，我们才开始接触 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.csdn.net/u013250416/article/details/78676784?locationNum=7&fps=1

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

从笔者电脑来看，数据集大小为11.2G，11000+张图片，训练速度还是挺慢的。训练速度较慢的原因有待研究。

04

—

模型测试

笔者认为“测试”的含义有两种，一种是利用数据集中的测试数据检测模型效果，叫test，另一种是利用数据集外实际应用中的数据集检测模型效果，叫deploy。以下分别介绍。

利用数据集中的测试集测试

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

./test.sh

开始测试，test.sh脚本的构造同train.sh，会调用solver_test.prototxt文件测试。测试结束后在屏幕打印detection_eval值和loss值。

利用数据集外的图片测试

这项测试的前提条件是，拥有属于自己的caffemodel。在caffe根目录/examples/MobileNet-SSD下运行

python merge_bn.py

即可生成指定名称的caffemodel。merge_bn.py文件中需要编辑的变量为：

train_proto ='/home/wluo/DeepLearning/CaffeLearning-weiliu89/caffe/examples/MobileNet-SSD/example/MobileNetSSD_train.prototxt' #train.prototxt路径

train_model = 'mobilenet_power_iter_25000.caffemodel' #should be your snapshot caffemodel

deploy_proto ='/home/wluo/DeepLearning/CaffeLearning-weiliu89/caffe/examples/MobileNet-SSD/example/MobileNetSSD_deploy.prototxt' #test.prototxt路径

save_model ='MobileNetSSD_deploy_20171226_iter25000.caffemodel' #指定生成的caffemodel名称

caffemodel生成之后，可以在caffe根目录/examples/MobileNet-SSD/执行

python demo.py

开始测试，在运行demo.py之前，可能需要做一些修改：

net_file='/home/wluo/DeepLearning/CaffeLearning-weiliu89/caffe/examples/MobileNet-SSD/example/MobileNetSSD_deploy.prototxt' #deploy.prototxt路径

caffe_model='MobileNetSSD_deploy_20171226_iter25000.caffemodel' #caffemodel名称

test_dir = "images_power" #测试图片文件夹名

CLASSES修改为自己数据集的类别和background。

如果一切顺利，运行demo.py后就可以看到每张测试图片的目标检测结果了，包括类别和概率。检测结果示例如下图所示。