《caffe学习之路》第六章：Ubuntu16.04 caffe ssd 训练自己的数据集

在写到这章时，博主已经身心疲惫。花了近一天的时间在探索如何利用caffe ssd训练自己的数据集，网上的资料太杂，写得太简要，我被绕得晕头转向，幸好找到了一篇简书：https://www.jianshu.com/p/2d342adc6654

里面说得真的是详细，有了它才有了今天这篇博客，写这篇博客只为自己做个记录，怕以后忘记了。所以想要知道如何利用caffe ssd训练自己的数据集，建议参考上面的这篇简书。

我先总结一下步骤：

一、制作 Pascal Voc 格式数据集：

第一步：找数据集，也就是你最后要识别什么类型的图片

第二步：用LabelImg工具对图片进行标记，每张图片都得到一个xml文件

第三步：写一个python小程序把图片分成训练集和测试集

二、利用caffe给的脚本把训练集和测试集生成lmdb文件

三、训练模型

下面对每个步骤进行详细说明：

1、制作 Pascal Voc 格式数据集

（1）建立 Pascal Voc 格式目录

cd ~/                      # 进入home目录
mkdir data                 # 新建一个data文件
cd data                    # 进入data目录    
mkdir VOCdevkit
cd VOCdevkit
mkdir MyDataSet            # 新建一个MyDataSet文件
cd MyDataSet
mkdir Annotations          # 存放 xml 文件
mkdir JPEGImages           # 存放 jpg 文件
mkdir ImageSets
mkdir lmdb                 # 可以不建，后面会自动生成           
cd ImageSets
mkdir Main

 参考图：

（2）标记数据

我懒得去网上下载图片，直接在VOC2012图片里面随便选择了34张人和25张车的图片，粘贴复制到JPEGImages目录下：

如果有强迫症，想有规律排列图片的序号，可以写一个脚本： 

#!/bin/bash
c=0                       # 从0开始排列
for i in *.jpg
do 
    mv -f $i $((c+=1)).jpg
done

（3）用LabelImg工具对图片进行标记

cd ~/                                                   # 进入home目录
mkdir git                                               # 创建一个git目录
git clone https://github.com/tzutalin/labelImg          # 下载

安装方法参考https://github.com/tzutalin/labelImg网址里面的教程，这里就不多说了。

最后把所有的xml文件存放到Annotations目录下。

参考图：

（3）划分训练集和测试集

下面是Python代码：

import os  
import random 
xmlfilepath=r'/home/bjw/data/VOCdevkit/MyDataSet/Annotations/'   #change xml path                        
saveBasePath=r"/home/bjw/data/VOCdevkit/"                        #change base path
trainval_percent=0.8                                             #adjust trainval percentage    
train_percent=0.8                                                #adjust train percentage        
total_xml = os.listdir(xmlfilepath)
num=len(total_xml)    
list=range(num)    
tv=int(num*trainval_percent)    
tr=int(tv*train_percent)    
trainval= random.sample(list,tv)    
train=random.sample(trainval,tr)    
  
print("train and val size",tv)  
print("traub suze",tr)  
ftrainval = open(os.path.join(saveBasePath,'MyDataSet/ImageSets/Main/trainval.txt'), 'w')    
ftest = open(os.path.join(saveBasePath,'MyDataSet/ImageSets/Main/test.txt'), 'w')    
ftrain = open(os.path.join(saveBasePath,'MyDataSet/ImageSets/Main/train.txt'), 'w')    
fval = open(os.path.join(saveBasePath,'MyDataSet/ImageSets/Main/val.txt'), 'w')    
  
for i  in list:    
    name=total_xml[i][:-4]+'\n'    
    if i in trainval:    
        ftrainval.write(name)    
        if i in train:    
            ftrain.write(name)    
        else:    
            fval.write(name)    
    else:    
        ftest.write(name)    
    
ftrainval.close()    
ftrain.close()    
fval.close()    
ftest .close()  

注意根据自己的路径进行修改，修改地方参考图： 

到这里Pascal Voc 格式数据集就算完成了。

2、生成LMDB文件

（1）首先在 Caffe 根目录下的 data 目录内创建一个名为 MyDataSet 的目录：

cd ~/caffe/data    #我的caffe目录在home目录下                          
mkdir MyDataSet  

然后执行下列命令：

cp VOC0712/create_* MyDataSet/                # 把create_list.sh和create_data.sh复制到MyDataSet目录                  
cp VOC0712/labelmap_voc.prototxt MyDataSet/   # 把labelmap_voc.prototxt复制到MyDataSet目录 

（2）修改create_list.sh

vim create_list.sh          # 修改 create_list.sh

修改如下： 

#!/bin/bash
# 如果你目录严格按照我上面提供的命令创建的话，那么下面 root_dir 等不用修改，直接用我的就行
# 如果你自定义了目录名，需要根据自己的定义修改
root_dir=$HOME/data/VOCdevkit/
sub_dir=ImageSets/Main
bash_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
for dataset in trainval test
do
  dst_file=$bash_dir/$dataset.txt
  if [ -f $dst_file ]
  then
    rm -f $dst_file
  fi
  for name in MyDataSet          # 如果建立的不是MyDataSet名目录，改成自己的
  do
# 注意下面这一段需要注释掉
#    if [[ $dataset == "test" && $name == "MyDataSet" ]]
#    then
#      continue
#    fi
    echo "Create list for $name $dataset..."
    dataset_file=$root_dir/$name/$sub_dir/$dataset.txt

    img_file=$bash_dir/$dataset"_img.txt"
    cp $dataset_file $img_file
    sed -i "s/^/$name\/JPEGImages\//g" $img_file
    sed -i "s/$/.jpg/g" $img_file

    label_file=$bash_dir/$dataset"_label.txt"
    cp $dataset_file $label_file
    sed -i "s/^/$name\/Annotations\//g" $label_file
    sed -i "s/$/.xml/g" $label_file

    paste -d' ' $img_file $label_file >> $dst_file

    rm -f $label_file
    rm -f $img_file
  done

  # Generate image name and size infomation.
  if [ $dataset == "test" ]
  then
    $bash_dir/../../build/tools/get_image_size $root_dir $dst_file $bash_dir/$dataset"_name_size.txt"
  fi

  # Shuffle trainval file.
  if [ $dataset == "trainval" ]
  then
    rand_file=$dst_file.random
    cat $dst_file | perl -MList::Util=shuffle -e 'print shuffle(<STDIN>);' > $rand_file
    mv $rand_file $dst_file
  fi
done

（3）修改create_data.sh

vim create_data.sh           # 修改 create_data.sh

修改如下： 

# 同样，如果你严格按照我的命令定义了目录名，就不需要修改
# 如果你修改了我上述命令中的目录名，需要你改的地方有 root_dir, data_root_dir, dataset_name，mapfile
cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir=$cur_dir/../..

cd $root_dir

redo=1
data_root_dir="$HOME/data/VOCdevkit"
dataset_name="MyDataSet"
mapfile="$root_dir/data/$dataset_name/labelmap_voc.prototxt"
anno_type="detection"
db="lmdb"
min_dim=0
max_dim=0
width=0
height=0

extra_cmd="--encode-type=jpg --encoded"
if [ $redo ]
then
  extra_cmd="$extra_cmd --redo"
fi
for subset in test trainval
do
  python $root_dir/scripts/create_annoset.py --anno-type=$anno_type --label-map-file=$mapfile --min-dim=$min_dim --max-dim=$max_dim --resize-width=$width --resize-height=$height --check-label $extra_cmd $data_root_dir $root_dir/data/$dataset_name/$subset.txt $data_root_dir/$dataset_name/$db/$dataset_name"_"$subset"_"$db examples/$dataset_name
done
# 修改好后 :wq 保存退出

（4）修改labelmap_voc.prototxt

vim labelmap_voc.prototxt    # 修改 label map

修改如下：

item {
  name: "none_of_the_above"
  label: 0
  display_name: "background"
}
item {
  name: "person"
  label: 1
  display_name: "person"
}
item {
  name: "car"
  label: 2
  display_name: "car"
}

 《注意》第一项background必须有，保持不变。后面的依次按照类别添加。

（5）生成LDBM数据格式

cd ~/caffe                               # 进入 Caffe 根目录
./data/MyDataSet/create_list.sh
./data/MyDataSet/create_data.sh

成功之后会在~/data/VOCdevkit/MyDataSet/lmdb和~/caffe/examples/MyDataSet目录下生成两个 lmdb文件

 到这里生成LMDB文件也完成了。

3、模型训练

（1）修改 ssd_pascal.py

终于进行到最后一步了 ，只需修改 ssd_pascal.py

cd ~/caffe/examples/ssd              # 进入 ssd_pascal.py 所在目录
cp ssd_pascal.py my_ssd_pascal.py    # 尽量不动 ssd_pascal.py 源文件
sudo gedit my_ssd_pascal.py          # 修改 my_ssd_pascal.py

修改如下：

# 82 行修改 LMDB 文件位置信息
# 上一步执行 *.sh 文件时候输出了 LMDB file 的存放位置
# The database file for training data. Created by data/VOC0712/create_data.sh
train_data = "examples/MyDataSet/MyDataSet_trainval_lmdb"
# The database file for testing data. Created by data/VOC0712/create_data.sh
test_data = "examples/MyDataSet/MyDataSet_test_lmdb"

# 258 行修改必要信息
# Stores the test image names and sizes. Created by data/VOC0712/create_list.sh
name_size_file = "data/MyDataSet/test_name_size.txt"
# The pretrained model. We use the Fully convolutional reduced (atrous) VGGNet.
pretrain_model = "models/VGGNet/VGG_ILSVRC_16_layers_fc_reduced.caffemodel"
# Stores LabelMapItem.
label_map_file = "data/MyDataSet/labelmap_voc.prototxt"
# MultiBoxLoss parameters.
num_classes = 3            # 修改为你要的分类数+1。例如我是2分类，我写了 2+1=3

# 332 行修改 GPU 信息
gpus = "0"   # 你要开启几个 GPU 加速就写几个，编号从0开始。我用一个2070 就写了 0

# 337 修改 batch size 大小
batch_size = 16    # 这个数字大小看你显存大小填写， 允许范围内越大越好

# 359 行修改测试集图片数
num_test_image = 12   # 根据你测试集图片数实际填写 
# 图片数为 $caffe_root/data/VOCdevkit/MyDataSet/ImageSets/Main/test.txt 的行数

# 修改好后 :wq 保存退出

（2）训练

cd ~/caffe
python example/ssd/my_ssd_pascal.py

模型存放位置：~/caffe/models/VGGNet/MyDataSet/MYSSD_300x300

参考图：

4、模型测试

测试代码在第五章，我做了一个图片测试，修改四个地方就可了：

第一处：file_type改为image

第二处：模型描述文件和caffemodel

第三处：检测图片的位置，找一个人和一个车的图片进行测试

第四处：lable

下面是检测结果： 

59张图片训练的效果看来还不错 。