【Caffe】用Caffe训练MobileNetSSD

一、制作训练数据集

具体过程请参考链接：【Caffe】制作自己的训练数据集。

二、编译Caffe

编译Caffe，能够训练跑通ssd。具体过程请参考链接：【Caffe】caffe-ssd训练自己的数据集。

三、训练MobileNetSSD

1. 克隆MobileNetSSD源码

git clone https://github.com/chuanqi305/MobileNet-SSD.git

将下载的整个文件夹拷贝到caffe/examples下面。

2. 生成网络文件

./gen_model.sh 31 #类别数+1

则在examples文件夹下生成了以下几个文件。

打开MobileNetSSD_train.prototxt，将第136行修改为自己的数据集label文件。

label_map_file: "labelmap3_testDataset.prototxt"

MobileNetSSD_test.prototxt需要做同样的修改。

3. 建立数据库链接

根据【Caffe】制作自己的训练数据集制作的数据集文件，这里做一个超链接。

ln -s /home/code/caffe/data/testDataSet/testDataSet_train_lmdb train_lmdb
ln -s /home/code/caffe/data/testDataSet/testDataSet_test_lmdb test_lmdb

就在当前目录多了两个超链接文件夹。

4. 修改超参数

打开solver_train.prototxt，将train_net和test_net修改成自己用gen_model.sh生成的MobileNetSSD_train.prototxt和MobileNetSSD_test.prototxt。test_iter修改为你的(测试图片数/batch size)。其他的超参数自己设置。

5. 训练

下载预训练权重文件。

打开train.sh，修改--weights为使用的预训练权重。执行

./train.sh

就开始训练了。然后在snapshot文件夹下就会生成一系列的.caffemodel和.solverstate文件了。

6. 合并bn层

为了提高模型运行速度，训练完成的caffemodel，将bn层合并到卷积层中。

打开merge_bn.py，修改caffe_root为自己的路径。

执行

python merge_bn.py --model examples/MobileNetSSD_deploy.prototxt --weights=snapshot/mobilenet_iter_100000.caffemodel

合并完成之后，在当前目录下多了no_bn.prototxt和no_bn.caffemodel两个文件。

7. 测试

打开demo.py文件，修改caffe_root为自己的路径。

修改net_file和caffe_model为要测试的文件。test_dir是存放测试图片的文件夹。

将CLASSES修改为自己数据集的类别名。

执行

python demo.py

终端打印出花费的时间。

同时会弹出窗口显示检测框。按空格键切换下一张图片。

检测效果不好，特别是小物体。没有tensorflow训练的检测效果好，可能是参数还没调好，需要进一步研究。