天天看點

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

這裡不再介紹 mmdetection 的安裝和配置,使用 mmdetection 較簡單的方法是使用已安裝 mmdetection 的 docker 容器。這樣直接省去了安裝 mmdetection 的過程,讓重心放在模型訓練上! 

如果你對 docker 和 mmdetection 還不是很熟悉,請自行搜尋一下,本文就不再贅述了。 

這裡附上 mmdetection 的 GitHub 位址: 

https://github.com/open-mmlab/mmdetection

0. 前期準備

首先預設你的電腦已經做好了下面這些前期準備工作: 

  • Ubuntu 16.04 或以上
  • GPU 安裝
  • cuda 安裝
  • cudnn 安裝
  • docker 安裝
  • nvidia-docker

當然,如果你連接配接的是公司或學校的伺服器,且伺服器已經做了上面幾點準備,那你隻需要一個 Xshell 遠端登入伺服器就行了。

下載下傳含 mmdetection 的 docker 鏡像

首先,我們需要找到一個已經配置好 mmdetection 環境的 docker 鏡像。可以在 dockerhub 上用 “mmdetection” 作為關鍵詞進行搜尋,也可以在 terminal 裡直接使用指令 docker search 進行搜尋:

$ docker search mmdetection      

結果顯示如下圖所示:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

這裡,我們選擇排第一的 vistart/mmdetection 鏡像,下載下傳 docker 鏡像的方法也很簡單,使用 docker pull 從鏡像倉庫中拉取指定鏡像:

$ docker pull vistart/mmdetection      

如果網絡沒問題,下載下傳會在幾分鐘之内完成。下載下傳完成之後,我們就可以檢視 vistart/mmdetection 鏡像是否已經放在本地鏡像裡了:

$ docker images      
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

可以看到 vistart/mmdetection 鏡像已經成功下載下傳了。

2. 建立含 mmdetection 的容器

包含 mmdetection 的鏡像已經下載下傳好了,下一步就是建立一個 docker 容器以供使用了:

$ docker run --runtime=nvidia --name mm_prj -i -t vistart/mmdetection /bin/bash      

對上面的指令解釋一下:--runtime=nvidia 很關鍵,能使建立的 docker 容器能使用主控端器的 GPU,不加這個參數則預設使用 CPU;--name mm_prj 是對建立的 docker 容器進行命名,該名稱為 mm_prj,讀者可自行修改。 

建立容器之後的界面如下:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

至此,名為 mm_prj 容器已經打開了。可以看到,該目錄中已經包含了 mmdetection 目錄,表示該 docker 鏡像已經安裝好了 mmdetection。

補充:

另外,補充一些退出容器、進入容器的操作。

退出容器:

# exit      

檢視現有容器:

$ docker ps -a      
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

可以看到,名為 mm_prj 的 docker 容器已經在容器清單了。 

打開容器:

$ docker start mm_prj
$ docker exec -i -t mm_prj /bin/bash      

3. 導入自己的 VOC 資料

這一步,我們需要把自己的資料打包成 Pascal VOC 格式。其目錄結構如下:

VOCdevkit

    --VOC2007

        ----Annotations

        ----ImageSets

            ------Main

        ----JEPGImages

簡單介紹一下,其中 Annotations 存放的是 .xml 檔案,JEPFImages 存放的是 .jpg 圖檔。 

按照此格式放置好自己的訓練資料之後,需要切分訓練資料和測試資料。在 VOCdevkit 目錄下建立一個 test.py 檔案。test.py 内容為:

import os
import random
trainval_percent = 0.8
train_percent = 0.8
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
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)
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('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()      

上面的代碼劃分資料集,trainval 占 80%,作為訓練集;test 占 20%,作為測試集。 

運作 test.py,将會在 VOCdevkit/ImageSets/Main 目錄下生成下面三個檔案:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

打開檔案可以看到,trainval.txt 包含訓練時所有的樣本索引,test.txt 包含測試時所有的樣本索引。 

自己的 VOC 資料制作完畢之後,從主控端(Ubuntu)複制到 /mmdetection/data/ 目錄下:

$ docker cp VOCdevkit mm_prj:/mmdetection/data/      

4. 修改 class_names.py 檔案

打開 /mmdetection/mmdet/core/evaluation/class_names.py 檔案,修改 voc_classes 為将要訓練的資料集的類别名稱。如果不改的話,最後測試的結果的名稱還會是’aeroplane’, ‘bicycle’, ‘bird’, ‘boat’,…這些。改完後如圖: 

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

5. 修改 voc.py 檔案

打開 mmdetection/mmdet/datasets/voc.py 檔案,修改 VOCDataset 的 CLASSES 為将要訓練的資料集的類别名稱。

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

如果隻有一個類,要加上一個逗号,否則将會報錯。

6. 修改配置檔案

mmdetection 中提供了很多目标檢測模型可供使用。例如,進入 /mmdetection/config/ 目錄,就會看到很多模型:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

根據我們選擇使用的模型,修改相應的配置檔案。本文我們使用的是FasterRCNN 模型,修改的是 faster_rcnn_r50_fpn_1x.py 檔案。 

6.1 修改 num_classes 變量 

打開 faster_rcnn_r50_fpn_1x.py,将 num_classes 變量改為:類别數 + 1(例如我有 20 類,是以改為 21):

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

6.2 修改 data_settings 

因為 faster_rcnn_r50_fpn_1x.py 預設使用的是 coco 資料集格式,我們要對其修改成相應的 VOC 資料格式。修改後的内容如下圖所示:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

6.3 調整學習率 

本文使用單 gpu 訓練,修改 img_per_gpu = 2,workers_per_gpu = 0。

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集
全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

對學習率的調整,一般遵循下面的習慣: 

  • 8 gpus、imgs_per_gpu = 2:lr = 0.02;
  • 2 gpus、imgs_per_gpu = 2 或 4 gpus、imgs_per_gpu = 1:lr = 0.005;
  • 4 gpus、imgs_per_gpu = 2:lr = 0.01

這裡,我們隻使用單 gpu,且 img_per_gpu = 2,則設定 lr = 0.00125。

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

這裡說一下 epoch 的選擇,預設 total_epoch = 12,learning_policy 中,step = [8,11]。total_peoch 可以自行修改,若 total_epoch = 50,則 learning_policy 中,step 也相應修改,例如 step = [38,48]。

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

至此,配置檔案已修改完畢。 

7. 模型訓練 

模型訓練非常簡單,隻需一行指令:

python3 ./tools/train.py ./configs/faster_rcnn_r50_fpn_1x.py      

注意執行上面的指令是在 /mmdetection 目錄下。 

如果有多個 gpu,例如 0, 1 号 gpu 都可用,則可以全部用起來訓練,指令如下:

CUDA_VISIBLE_DEVICES=0,1 python3 ./tools/train.py ./configs/faster_rcnn_r50_fpn_1x.py --gpus 2      

上面的 --gpus 2 表示使用的 gpu 個數為 2。如果使用多塊 gpu,注意修改學習率 lr。

然後,訓練就開始了:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

從列印出的資訊中,可以看到目前的 epoch 和 loss 值。 

每個 epoch 會生成一個模型,并自動儲存在 /mmdetection/work_dirs/faster_rcnn_r50_fpn_1x/ 目錄下。

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

訓練完成之後,latest.pth 即 epoch_12.pth 就是最終的模型。 

8. 模型測試,計算 mAP 

下面我們将使用訓練好的模型對測試集進行驗證,并計算 mAP。 

8.1 生成 pkl 檔案 

首先,生成 pkl 檔案:

python3 ./tools/test.py ./configs/faster_rcnn_r50_fpn_1x.py ./work_dirs/faster_rcnn_r50_fpn_1x/latest.pth --out=result.pkl      

8.2 計算測試集 

mAP 對測試集計算 mAP,隻需一行指令:

python3 ./tools/voc_eval.py result.pkl ./configs/faster_rcnn_r50_fpn_1x.py      

計算結果如下:

全網最細 | 教你如何在 docker 容器下使用 mmdetection 訓練自己的資料集

圖中可以看到,最後計算的 mAP = 0.978。(本文中的目标檢測場景比較簡單,目标清晰明确,故 mAP 很高)

9. 總結

好了,以上就是教你如何一步一步在 docker 容器中使用 mmdetection 來訓練自己的資料集并測試。建議大家使用自己的資料集嘗試跑一下,看下效果~