目標檢測算法SSD之訓練自己的數據集

prerequesties 預備知識/前提條件

下載和配置了最新SSD代碼

git clone https://github.com/weiliu89/caffe ~/work/ssd
cd $_
git checkout ssd

編譯caffe

下載必要的模型（包括prototxt和caffemodel）；

運行了evaluation和webcam的例子，會提示caffe的import報錯。添加pycaffe路徑到PYTHONPATH環境變量，或者寫一個_init_paths.py來輔助引入都可以（推薦后者）。

准備自己的數據集

做成VOC2007格式的：

JPEGImages/*.png
ImageSets/Main/*.txt
Annotations/*.xml

這3個目錄

生成訓練用的lmdb數據

我這里數據集名叫traffic_sign，放在/home/chris/data/traffic_sign

1.復制原有腳本文件

cd ~/work/ssd
cp -R data/VOC0712 data/traffic_sign

2.修改data/traffic_sign/create_list.sh

#!/bin/bash

#root_dir=$HOME/data/VOCdevkit/
root_dir=$HOME/data/
sub_dir=ImageSets/Main
bash_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" #當前文件所在目錄
for dataset in train test
do
  dst_file=$bash_dir/$dataset.txt
  if [ -f $dst_file ]
  then
    rm -f $dst_file
  fi
  for name in traffic_sign
  do
    if [[ $dataset == "test" && $name == "VOC2012" ]]
    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/$/.png/g" $img_file     #在行尾追加.png后綴

    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   #在行尾追加.xml后綴

    paste -d' ' $img_file $label_file >> $dst_file  #img_file和label文件的對應行拼接

    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 train file.
  if [ $dataset == "train" ]
  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.修改data/traffic_sign/create_data.sh：

#!/bin/bash

cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir=$cur_dir/../..

cd $root_dir

redo=1
data_root_dir="$HOME/data"
#dataset_name="VOC0712"
dataset_name="traffic_sign"
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=png --encoded"
if [ $redo ]
then
  extra_cmd="$extra_cmd --redo"
fi
for subset in test train
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

4.修改data/traffic_sign/labelmap_voc.prototxt

item {
  name: "none_of_the_above"
  label: 0
  display_name: "background"
}
item {
  name: "sign"
  label: 1
  display_name: "sign"
}

5.生成數據

# 確保你還是在ssd代碼根目錄，比如我是~/work/ssd
./data/traffic_sign/create_list.sh
./data/traffic_sign/create_data.sh

執行訓練

依然需要修改ssd默認的訓練腳本內容，來匹配自己的數據集。

1.復制原有訓練腳本

cd ~/work/ssd
cd examples/ssd
cp ssd_pascal.py ssd_traffic.py

2.修改訓練腳本

編輯ssd_traffic.py內容，修改：

• 數據集指向
  train_data和test_data , 指向examples中你的數據，例如:

train_data = "examples/traffic_sign/traffic_sign_train_lmdb"
test_data = "examples/traffic_sign/trainffic_sign_test_lmdb"

這里很奇怪，我的examples/traffic_sign/目錄下確實有這兩個lmdb的文件夾，是指向~/data/traffic_sign/lmdb/目錄下的兩個lmdb文件夾，但是訓練時提示lmdb錯誤。

換成鏈接文件的源文件，也就是寫絕對路徑，就不報錯了。

• 測試圖像數量

num_test_image 該變量修改成自己數據集中測試數據圖片的數量

• 類別數

num_classes 該變量修改成自己數據集中 標簽類別數量數 + 1

• gpu選項

gpus = "0,1,2,3" 電腦有幾個gpu就寫多少個，如果有一個就寫gpus="0",兩個就寫gpus="0,1",以此類推

• 迭代次數

solver_param = {
	...
	'stepvalue': [50000, 60000, 70000],
	'max_iter': 70000,
	'snapshot': 10000,
}

• 各種VOC0712換成自己數據集的名字（我的是traffic_sign)

model_name = "VGG_traffic_sign_{}".format(job_name)
save_dir = "models/VGGNet/traffic_sign/{}".format(job_name)
snapshot_dir = 
job_dir = 
name_size_file = 
label_map_file = 

• batch_size
  比如6G顯存的970顯卡，跑不起來SSD。修改：

batch_size = 16   # 32->16
accum_batch_size = 16  # 32->16

此時顯存占用為4975MiB

如果你顯存很大，與其閑置不如使用它，調大batch_size即可

• base_lr
  調整了batch_size或單純因為數據集的原因，導致出現loss為nan的情況，考慮減小學習率，這里通過減小base_lr實現。

3.執行訓練

cd ~/work/ssd    #務必到ssd的根目錄執行
python examples/ssd/ssd_traffic.py

## reference

https://my.oschina.net/u/1046919/blog/777470