基于自制数据集的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文件,该文件中的类别为“准备数据集”部分的类别数+1background)。笔者的数据集共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后就可以看到每张测试图片的目标检测结果了,包括类别和概率。检测结果示例如下图所示。

  • 发表于:
  • 原文链接http://kuaibao.qq.com/s/20171230G0BQC800?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 yunjia_community@tencent.com 删除。

扫码关注云+社区

领取腾讯云代金券