Caffe-SSD-Ubuntu16-04-訓練自己的資料集

catn1p發表於2019-03-03
Ubuntu

總的來說，Caffe 是一個比較難上手的框架。這次嘗試訓練 Caffe 框架下 SSD 模型的訓練是我第一次使用 Caffe 框架。下面就說一說我踩過的幾個坑，希望能夠幫助到大家。

1 編譯 Caffe 框架

這一步是我認為使用 Caffe 框架的最大障礙，編譯不停出錯。最後我不得不放棄轉而使用 Docker 解決 Caffe 的編譯安裝問題。下面寫出 Docker 的安裝以及拉取所需映象的方法。

Docker 安裝

安裝過程我是參照的 Docker 官方的安裝指引，傳送門在這裡。
我使用的安裝命令為(Ubuntu 16.04 LTS)：

# If you have installed older version of docker, removing it by using command as follows
sudo apt-get remove docker docker-engine docker.io

# Docker installation
# 1. Update the apt package index
sudo apt-get update
# 2. Install packages to allow apt to use a repository over HTTPS
sudo apt-get install 
    apt-transport-https 
    ca-certificates 
    curl 
    software-properties-common
# 3. Add Docker’s official GPG key
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
# 4. Add apt repository
sudo add-apt-repository 
   "deb [arch=amd64] https://download.docker.com/linux/ubuntu 
   $(lsb_release -cs) 
   stable"
# 5. INSTALL DOCKER CE
sudo apt-get update
sudo apt-get install docker-ce
# 6. Verify that Docker CE is installed correctly by running the hello-world image
sudo docker run hello-world
複製程式碼

以上就是 Docker 的安裝過程，安裝完成 Docker 後，下一步就是尋找合適的映象並拉取。如果你對 Docker 的使用並不熟悉，推薦你閱讀《第一本Docker書》的前4章，閱讀時間大約在2個小時，下載在這裡。

尋找並拉取合適的映象

我是在 Docker Hub 直接搜尋我需要的映象的，我的需求是 python2, gpu 版本的 Caffe 並且是 SSD 分支。
在 Bing.com 搜尋「docker hub」，開啟之後搜尋關鍵詞 「caffe ssd」。

搜尋映象

我所使用的映象連結在這裡。
如果你想在 Docker 中使用 GPU 加速，那麼你還必須安裝 nvidia-docker，不過幸好安裝特別簡單。過程如下：

# If you have nvidia-docker 1.0 installed: we need to remove it and all existing GPU containers
docker volume ls -q -f driver=nvidia-docker | xargs -r -I{} -n1 docker ps -q -a -f volume={} | xargs -r docker rm -f
sudo apt-get purge -y nvidia-docker

# Add the package repositories
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | 
  sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | 
  sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update

# Install nvidia-docker2 and reload the Docker daemon configuration
sudo apt-get install -y nvidia-docker2
sudo pkill -SIGHUP dockerd

# Test nvidia-smi with the latest official CUDA image
docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi
複製程式碼

按照順序依次執行就可以了，如果你遇到問題，請參照 nvidia-docker 安裝官方指引。

拉取映象：

sudo nvidia-docker run -it --name xxx -v /your/path/to/swap place(host):/your/path/to/swap place narumi/caffe-ssd-gpu /bin/bash
# -it 指定這個 Docker 容器是可互動的，不可少
# --name 指定 Docker 容器的名稱，方便以後使用 。將 xxx 替換為你想要的名稱，例如我指定的名稱為 caffe_ssd_gpu_py2。
# -v 指定掛載目錄到容器中，這樣方便容器與宿主機進行檔案交換。 「:」前為宿主機目錄，後為容器內目錄
# narumi/caffe-ssd-gpu 指定我拉取的映象名
# /bin/bash 命令使得容器開啟後自動為我開啟終端

# 我自己使用的命令如下，供大家參考
sudo nvidia-docker run -it --name caffe_ssd_gpu_py2 -v /home/ubuntu/work/docker_swap:/home/swap narumi/caffe_ssd_gpu /bin/bash
# 這樣，我在容器中訪問 /home/swap 時就能看到我主機 /home/ubuntu/work/docker_swap 下存放的檔案了
複製程式碼

2 製作自己的資料集

這一步才是重點，下面我們將處理自己的使用的資料集。

標註資料集

因為我們將要把資料集製作為 pascal voc 格式的資料集，因此我們需要將標註資訊存放在 xml 檔案中。推薦使用 LabelImg 這個開源軟體標註。其優點是圖形化介面，並且自動生成 xml 檔案，省去了很多轉化的步驟，傳送門在這裡。
怎樣使用其實非常簡單，在這裡略過不講。如果你不會使用，請善用搜尋。無法解決的話可以留言詢問。

製作 Pascal Voc 格式資料集

在主機的交換目錄下建立名為 VOCdevkit 的資料夾用於存放圖片等內容。

交換目錄就是剛剛建立容器時候指定的主機目錄，例如我的交換目錄為 /home/ubuntu/work/docker_swap

具體命令如下：

cd /home/ubuntu/work/docker_swap
mkdir VOCdevkit
cd VOCdevkit
mkdir VOC2007
cd VOC2007
mkdir Annotations          # 存放 xml 檔案
mkdir JPEGImages         # 存放 jpg 檔案
mkdir ImageSets
cd ImageSets
mkdir Main
複製程式碼

建立完成後將 jpg 和 xml 檔案放入對應目錄下。然後使用 python 指令碼劃分一下訓練集和測試集。我把我使用的指令碼貼在這裡：

import os  
import random   

# 下面兩個目錄改成自己的目錄  
xmlfilepath=r`/your/path/to/xmls`                            
saveBasePath=r"your/path/to/save/VOCdevkit"  
  
trainval_percent=0.9           # 劃分訓練集和驗證集的比例
train_percent=0.9               # trainval 中 訓練集所佔比例
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,`VOC2007/ImageSets/Main/trainval.txt`), `w`)    
ftest = open(os.path.join(saveBasePath,`VOC2007/ImageSets/Main/test.txt`), `w`)    
ftrain = open(os.path.join(saveBasePath,`VOC2007/ImageSets/Main/train.txt`), `w`)    
fval = open(os.path.join(saveBasePath,`VOC2007/ImageSets/Main/val.txt`), `w`)    
  
for i  in list:    
    name=total_xml[i][:-4]+`
`    
    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()  
複製程式碼

執行完畢後 VOCdevkit/VOC2007/ImageSets 下應該有4個 txt 文字。

修改 create_list.sh 與 create_data.sh 並生成 LMDB files

首先在 Caffe 根目錄下的 data 目錄內建立一個名為 VOC2007 的目錄，然後執行下列命令：

cd /opt/caffe/data
mkdir VOC2007
cp VOC0712/create_* VOC2007/
cp VOC0712/labelmap_voc.prototxt VOC2007/
cd VOC2007
# 修改 label map
vim labelmap_voc.prototxt
# 如果提示沒有vim，使用 sudo apt-get install vim 安裝一下
複製程式碼
# labelmap_voc.prototxt 中內容修改為自己需要的內容
item {
  name: "none_of_the_above"
  label: 0
  display_name: "background"
}
item {
  name: "label1"                    # label 為你自己資料集裡label的名稱，替換即可
  label: 1
  display_name: "label1"
}
item {
  name: "label2"
  label: 2
  display_name: "label2"
}
item {
  name: "label3"
  label: 3
  display_name: "label3"
}
item {
  name: "label4"
  label: 4
  display_name: "label4"
}
...
# 修改好後 :wq 儲存
複製程式碼
# 修改 create_list.sh
vim create_list.sh 
複製程式碼
#!/bin/bash
# 如果你目錄嚴格按照我上面提供的命令建立的話，那麼下面 root_dir 等不用修改，直接用我的就行
# 如果你自定義了目錄名，需要根據自己的定義修改
root_dir=/opt/caffe/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 VOC2007
  do
# 注意下面這一段需要註釋掉
#    if [[ $dataset == "test" && $name == "VOC2007" ]]
#    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
# 修改好 :wq 儲存退出
複製程式碼
# 修改 create_data.sh
vim create_data.sh
複製程式碼
# 同樣，如果你嚴格按照我的命令定義了目錄名，就不需要修改
# 如果你修改了我上述命令中的目錄名，需要你改的地方有 root_dir, data_root_dir, dataset_name, mapfile
cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir="/opt/caffe"

cd $root_dir

redo=1
data_root_dir="/opt/caffe/data/VOCdevkit"
dataset_name="VOC2007"
mapfile="/opt/caffe/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 儲存退出
複製程式碼

都修改好後執行指令碼生成 LMDB

# 進入 Caffe 根目錄
cd /opt/caffe
./data/VOC2007/create_list.sh
./data/VOC2007/create_data.sh
複製程式碼

3 修改 ssd_pascal.py 並開始訓練

終於進行到最後一步了，這一步相對來說很簡單。

# Caffe root dir
cd /opt/caffe
vim example/ssd/ssd_pascal.py
複製程式碼
# 82 行修改 LMDB 檔案位置資訊
# 上一步執行 *.sh 檔案時候輸出了 LMDB file 的存放位置
# The database file for training data. Created by data/VOC0712/create_data.sh
train_data = "/opt/caffe/data/VOCdevkit/VOC2007/lmdb/VOC2007_trainval_lmdb"
# The database file for testing data. Created by data/VOC0712/create_data.sh
test_data = "/opt/caffe/data/VOCdevkit/VOC2007/lmdb/VOC2007_test_lmdb"

# 258 行修改必要資訊
# Stores the test image names and sizes. Created by data/VOC0712/create_list.sh
name_size_file = "data/VOC2007/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/VOC2007/labelmap_voc.prototxt"
# MultiBoxLoss parameters.
num_classes = 5            # 修改為你要的分類數+1。例如我是4分類，我寫了 4+1=5

# 332 行修改 GPU 資訊
gpus = "0,1"   # 你要開啟幾個 GPU 加速就寫幾個，編號從0開始。我用兩張 1080TI 就寫了 0,1

# 337 修改 batch size 大小
batch_size = 32    # 這個數字大小看你視訊記憶體大小填寫， 允許範圍內越大越好

# 359 行修改測試集圖片數
num_test_image = 1000   # 根據你測試集圖片數實際填寫 
# 圖片數為 $caffe_root/data/VOCdevkit/VOC2007/ImageSets/Main/test.txt 的行數

# 修改好後 :wq 儲存退出
複製程式碼

下一步就是訓練，但是開始之前，先下載預訓練的 base model 放入對應位置，減少訓練時間並提高效率
為了方便大家下載，我上傳到百度網盤了，密碼: ip6v。如果你想自己下載，去 github 下載，連結。

 # 下載好的 caffe model 放在一開始指定的交換區
cd /opt/caffe/model
mkdir VGGNet
cp /home/swap/VGG_ILSVRC_16_layers_fc_reduced.caffemodel /ope/caffe/model/
複製程式碼

開始訓練

cd /opt/caffe
python example/ssd/ssd_pascal.py
複製程式碼
Train Log

結語

至此，恭喜你可以使用 Caffe 框架訓練屬於自己的 SSD 模型了， SSD 的官方實現的 github 地址為 連結。如果你有任何問題，可以求助於 github 討論區或者留言問詢。

另外，歡迎關注我。最近會寫很多關於目標檢測方面的文章，包括一些論文翻譯以及論文解讀。還可能寫一點演算法實現的內容。

謝謝閱讀。

相關文章