YOLO 簡單訓練自己的資料集

tags: tutorials YOLO Deep Learning

使用工具

作業系統

• Ubuntu 18.04

GPU

• Nvidia GTX1060 6GB

安裝 YOLO Darknet 環境與測試

安裝YOLO

1. 下載 YOLO 官方的程式碼

$ git clone https://github.com/pjreddie/darknet

2. 進入該資料集

$ cd darknet

3. (Optional) GPU 版本需修改編譯設定檔 Makefile，將 GPU 與 CUDNN 設定為 1

$ vim Makefile

GPU=1
CUDNN=1

4. 開始編譯

$ make

測試

1. 下載預訓練好的權重

wget https://pjreddie.com/media/files/yolov3.weights

2. 測試辨識，即會得到一張預測結果圖 predictions.jpg

$ ./darknet detect cfg/yolov3.cfg yolov3.weights data/dog.jpg

訓練自己的資料集（以分類汽車、機車、行人, 訓練yolov3-tiny為例）

將資料轉換成 YOLO 的格式

---

圖片資料夾

• image
  • 00001.jpg
  • 00002.jpg
    .
    .
  • 08000.jpg

標籤資料夾

• label
  • 00001.xml
  • 00002.xml
    .
    .
  • 08000.xml

下載voc2yolo.py，此檔案可將voc格式轉換成yolo格式（下載連結）
修改voc2yolo.py裡的圖片與標籤資料夾的路徑

xmlFolder = "label"    # voc的xml檔位置
imgFolder = "image"  # img檔位置
imgYoloFolder = "yolo"    # 轉換後img和txt儲存位置
classList = {"vehicle":0, "pedestrian":1, "rider":2}    # 類別的名稱

運行voc2yolo.py將voc的格式轉換成YOLO的格式，得到以下的資料夾

YOLO格式的標籤與圖片的資料夾

• yolo
  • 00001.jpg
  • 00001.txt
  • 00002.jpg
  • 00002.txt
    .
    .
  • 08000.jpg
  • 08000.txt

準備訓練資料

---

資料夾架構

• cfg_own_data
  • obj.data
  • obj.names
  • train.txt
  • valid.txt
  • yolov3-tiny.cfg

1. obj.data：設定類別數目、訓練與驗證圖片位置的紀錄檔、obj.names的位置、儲存權重的資料夾

classes= 3
train = cfg_own_data/train.txt
valid = cfg_own_data/valid.txt
names = cfg_own_data/obj.names
backup = weights

2. obj.names：設定類別的名稱

vehicle
pedestrian
rider

下載write_train_valid_txt.py，此檔案可將資料集隨機打散成訓練集和測試集（下載連結）
修改write_train_valid_txt.py的圖片路徑

imgfolder = "image"    # 圖片路徑
yolofolder = "yolo"    # yolo資料夾的路徑
train_txt_path = "cfg_own_data/train.txt"    # 儲存train.txt的路徑
valid_txt_path = "cfg_own_data/valid.txt"    # 儲存valid.txt的路徑

運行write_train_valid_txt.py即可隨機打亂資料分成train.txt與valid.txt

3. train.txt

yolo/05839.jpg
yolo/03954.jpg
yolo/02943.jpg
.
.

4. valid.txt

yolo/01273.jpg
yolo/07432.jpg
yolo/01234.jpg
.
.

5. yolov3-tiny.cfg
   從 cfg 資料夾中複製 yolov3-tiny.cfg 到這個資料夾中，並修改訓練設定
   1.
   ​​​​batch=32
   ​​​​subdivisions=8
   

   2. 第一個 Detector，filters 通常設為
      $(classes+5)\ast 3$
   ​​​​filters=24
   

   ​​​​classes=3
   

   3. 第二個 Detector
   ​​​​filters=24
   

   ​​​​classes=3
   

開始訓練

---

在 darknet 資料夾下，創一個資料夾放訓練時輸出的權重檔

$ mkdir weights

下載在ImageNet預訓練好的模型 darknet53.conv.74

$ wget https://pjreddie.com/media/files/darknet53.conv.74

開始訓練

$ ./darknet detector train cfg_own_data/obj.data cfg_own_data/yolov3-tiny.cfg 

每 100 個 batch 會輸出一次 weight檔，下圖為訓練過程的 log

若圖片數量至上千張，average loss error 達 0.06 即可停止；
若圖片數量只有幾百張，average loss error 達 0.6 即可停止

測試訓練結果

---

$ ./darknet detector test cfg_own_data/obj.data cfg_own_data/yolov3-tiny.cfg weight/yolov3-tiny_10000.weights