在 Colab 上利用 Yolov3 框架和自有標註資料來訓練自己的物件辨識系統

訓練成果

為什麼會有這篇文章

其實原意是利用圖片辨識手遊的畫面，來幫我自動玩手遊練功。在找資料的時候，很多人都推薦 Yolov3，優點是非常快速，剛好利用在畫面變動巨大的手遊上。

一開始入門的時候，光是如何在 Colab 引入 Yolo、如何把資料轉換成 Yolo 格式，和如何訓練就花了我數天；大多文章都只講到片段，沒有整合的版本。希望這篇文章能幫助到後來的人。

這篇文章會教你

在此文章中，我們將會訓練一個能偵測圖片中貓和狗的模型。你可以自由替換資料集，還偵測自訂的物體。

除此之外，還有：

• 利用 Colab 128G RAM GPU 來訓練你的 Yolo3 模型
• 掛載 Google Drive 檔案到 Colab 檔案系統中
• 將 PASCAL VOC 標籤格式轉換成 Yolo 用的標籤格式
• 產生 Yolo 訓練需要的 cfg 設定檔案
• 將訓練後的 weight 檔案同步至 Google Drive 中，避免遺失
• 如何利用 weight 檔案來辨識圖片中的內容

名詞解釋：Darknet 和 Yolo 是什麼關係

這邊借用 Tommy Huang 的解釋：

YOLO 是 Joseph Redmon 提出的演算法。Darknet 是 YOLO 作者(Joseph Redmon)自己寫出來的 deep learning framework，可以想像 darknet是另一個tensorflow，但功能沒這麼強大，主要是用來實現 YOLO系列的算法，但也是可以用來實現其它算法(效果可能沒這麼好)。

在實際應用中，把他們當作同一個東西的不同名字就可以了。

事前準備：編譯 Darknet 執行檔案

在開始前，你需要先編譯好 Darknet 執行檔案。你可以參考我的另外一篇文章「如何在 Colab 安裝 Darknet 框架訓練 YOLO v3 物件辨識並且最佳化 Colab 的訓練流程」，文章中會將編譯好的 Darknet 執行檔案放到 Google Drive目錄下。本文章將會利用這個執行檔案進行訓練。

完整 Colab 範例程式碼請見 http://bit.ly/2XCOT3E

步驟一：將 Google Drive 掛載到 Colab 目錄下

from google.colab import drive
drive.mount(‘/drive’, force_remount=True)

執行後，在執行結果的地方會出現一串授權網址。進入網址後，選擇要掛載 Colab 的帳號後，會出現一串授權碼。請把那串授權碼貼入輸入框中。如果沒有出現任何錯誤，那就是掛載成功了。

點選網址取得授權碼，然後再把授權碼貼回下面的輸入框中

步驟二：準備訓練資料

在此範例中，我們用 The Oxford-IIIT Pet Dataset 這個資料集來當作示範，裡面含有 7000 多張貓和狗的圖片。我們將利用此資料集來訓練一個可以偵測貓和狗的模型。

準備好的資料將會放在 Colab 檔案系統中的 /content/pet_detection 中。

# download the example dataset
!wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
!wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz
# move image and label folder into pet_detection folder
!mkdir /content/pet_detection
!tar -xf images.tar.gz -C /content/pet_detection
!tar -xf annotations.tar.gz
!mv annotations/xmls /content/pet_detection/labels
!rm -fr annotations

如果你要使用自己的資料集，可以把圖片檔案放在 /content/pet_detection/images，把標記 XML 檔案放在 /content/pet_detection/labels 中。

如果你想要手動標記，推薦使用 labelImg 這套程式。可以參考「https://blog.gtwang.org/useful-tools/labelimg-graphical-image-annotation-tool-tutorial/」

步驟三：節取出所有標籤的名字

由於 darknet 框架會將物體名字全部轉成數字，我們需要先將物體名字全部擷取出來存在一份檔案中，當作之後的對照表。

import glob
import os
import re
labels = set()
for path in glob.glob(os.path.join(LOCAL_LABELS_DIR_PATH, "*.xml")):
    with open(path, 'r') as f:
        content = f.read()
    # extract label names
    matches = re.findall(r'<name>([\w_]+)<\/name>', content, flags=0)
    labels.update(matches)
# write label into file
with open(os.path.join(LOCAL_CFG_DIR_PATH, "obj.names"), 'w') as f:
    f.write("\n".join(labels))
print('Read in %d labels: %s' % (len(labels), ", ".join(labels)))

步驟四：將 PASCAL VOC 標記資料轉換成 YOLO 格式的標記資料

Yolo 不是使用標準的格式，原本的 VOC 標記格式需要轉換後才能使用在 darkent 框架上。

這邊就不詳細解釋如何轉換，對如何轉換的詳細規格可以參考 Yolo 官網 。我們直接使用我從 convert2Yolo 套件中擷取出來的片段程式碼來執行轉換，並把轉換的結果都放到 /content/pet_detection/yolos 目錄中。

import sys
sys.path.append('/content/app')
from Format import VOC, YOLO
voc = VOC()
yolo = YOLO(os.path.join(LOCAL_CFG_DIR_PATH, "obj.names"))
flag, data = voc.parse(LOCAL_LABELS_DIR_PATH)
flag, data = yolo.generate(data)
flag, data = yolo.save(data,
    save_path=LOCAL_YOLOS_DIR_PATH,
    img_path=LOCAL_IMAGES_DIR_PATH, img_type=".jpg", manipast_path="./")

步驟五：準備訓練用的設定檔

Darknet 訓練時，總共需要以下四個檔案：

• obj.names：所有的物體標籤名稱，每一行一個。例如此例中只會有兩行，分別是 cat dog 。
• yolov3.cfg：darknet 網路的設定檔，描述每一層網路應該要如何建立，以及建立多少 node 等。裡面有些數值需要根據你的訓練資料來個別設定。
• train.txt test.txt ：這兩個檔案告訴 darknet 要到哪個路徑下找到訓練用的圖片。
• obj.data：darknet 的主要設定檔案，告訴 darknet 其他的設定檔路徑。darknet 會一一去讀取其他的檔案。

首先我們先產生 obj.names 檔案以及 yolov3.cfg檔案

# create the cfg file
!cp /content/app/darknet_cfg/yolov3.cfg {LOCAL_CFG_DIR_PATH}/yolov3.cfg
# fetch label_names
with open(os.path.join(LOCAL_CFG_DIR_PATH, "obj.names"), 'r') as f:
  f_content = f.read()
label_names = f_content.strip().splitlines()
# update the cfg file
with open(os.path.join(LOCAL_CFG_DIR_PATH, "yolov3.cfg"), 'r') as f:
  content = f.read()
with open(os.path.join(LOCAL_CFG_DIR_PATH, "yolov3.cfg"), 'w') as f:
  num_max_batches = len(label_names)*2000
  content = content.replace("%NUM_CLASSES%", str(len(label_names)))
  content = content.replace("%NUM_MAX_BATCHES%", str(num_max_batches))
  content = content.replace("%NUM_MAX_BATCHES_80%", str(int(num_max_batches*0.8)))
  content = content.replace("%NUM_MAX_BATCHES_90%", str(int(num_max_batches*0.9)))
  content = content.replace("%NUM_CONVOLUTIONAL_FILTERS%", str((len(label_names)+5)*3))
f.write(content)
!cp {LOCAL_CFG_DIR_PATH}/obj.names {GDRIVE_CFG_DIR_PATH}
!cp {LOCAL_CFG_DIR_PATH}/yolov3.cfg {GDRIVE_CFG_DIR_PATH}

接著讀取圖片，產生 train.txt test.txt兩個標示訓練用圖片路徑的檔案

# create train and test files
import random
import glob
txt_paths = glob.glob(os.path.join(LOCAL_YOLOS_DIR_PATH, "*.txt"))
random.shuffle(txt_paths)
num_train_images = int(len(txt_paths)*0.8)
with open(os.path.join(LOCAL_CFG_DIR_PATH, "train.txt"), 'w') as f:
  for path in txt_paths[:num_train_images]:
    f.write("%s/%s\n" % (LOCAL_YOLOS_DIR_PATH, os.path.basename(path).replace(".txt", ".jpg")))
with open(os.path.join(LOCAL_CFG_DIR_PATH, "test.txt"), 'w') as f:
  for path in txt_paths[num_train_images:]:
    f.write("%s/%s\n" % (LOCAL_YOLOS_DIR_PATH, os.path.basename(path).replace(".txt", ".jpg"))

然後產生最後的 obj.data 檔案

# create obj
with open(os.path.join(LOCAL_CFG_DIR_PATH, "obj.data"), 'w') as f:
  f.write("classes=%d\n" % (len(label_names)))
  f.write("train=%s/train.txt\n" % (LOCAL_CFG_DIR_PATH))
  f.write("valid=%s/test.txt\n" % (LOCAL_CFG_DIR_PATH))
  f.write("names=%s/obj.names\n" % (LOCAL_CFG_DIR_PATH))
  f.write("backup=%s\n" % (GDRIVE_WEIGHTS_DIR_PATH))

步驟六：準備 darkent 執行檔

我們直接從之前已經編譯好的檔案複製過來就好，不用每次都重頭編譯那實在是太~花~時~間~了~。編譯的方法請見 「如何在 Colab 安裝 Darknet 框架訓練 YOLO v3 物件辨識並且最佳化 Colab 的訓練流程」這邊文章。

# copy the pretrained darknet bin file
!cp {GDRIVE_DARKNET_BIN_FILE_PATH} /content/
!chmod +x /content/darknet

步驟七：（可選）使用 darknet 預先訓練的基底模型

Darknet 也好心的提供了預先訓練的模型，以此為基底，可以讓後來的訓練比較快達到較好的辨識率。但前提是你的圖片都是常見的圖片，例如一般照片、場景照片等；如果是一些遊戲畫面很少見的，從 0 開始訓練可能會達到比較好的效果。

# Use the pre-trained weights to speed up the training speed
!wget https://pjreddie.com/media/files/darknet53.conv.74

步驟八：開始訓練模型

終於準備好所有的資料了，可以開始訓練模型。回顧一下，所有需要的東西共有這些檔案

content
├── pet_detection
│   ├── cfg
│   │   ├── obj.data
│   │   ├── obj.names
│   │   ├── test.txt
│   │   ├── train.txt
│   │   └── yolov3.cfg
│   └── yolos
│
├── darknet
└── darknet53.conv.74

訓練的指令是：

# train the model
!./darknet detector train {LOCAL_CFG_DIR_PATH}/obj.data {LOCAL_CFG_DIR_PATH}/yolov3.cfg darknet53.conv.74

訓練過程中會顯示出訓練階段的訊息，例如：

可以從 avg loss 來判斷訓練成果，據網路上大家的說法，到 0.6 左右效果就是不錯了。我自己的經驗是，還是看資料集，通常一開始都是 1000 ~ 2000 左右，後來會降到 20~1 左右。可以隨時拿訓練模型的 weights 來試試看成果來決定要不要繼續訓練。

小技巧：如何從上次訓練的 weights 繼續往下訓練

由於我們是在 Colab 訓練，常常會遇到斷線需要重新開始。所以有個小技巧可以直接從上次訓練的地方繼續往下訓練，就不用每次都重頭。

只要把原本指令中的 darknet53.conv.74換成最後儲存下來的 weights，就可以接續訓練下去。

# train the model
!./darknet detector train {LOCAL_CFG_DIR_PATH}/obj.data {LOCAL_CFG_DIR_PATH}/yolov3.cfg {GDRIVE_WEIGHTS_DIR_PATH}/yolov3_last.weights

如果成功的話，你就可以看到 avg loss 不是從 1000~2000 開始跳，而是從上次的最後數值開始的。

小技巧：如何不要顯示那麼多執行 log，避免瀏覽器當機

由於我們是在 Colab 訓練，如果訓練的時間一長，太多的 Log 會讓瀏覽器效能變得非常差甚至當機。所以有個小技巧，可以讓 darknet 不要噴出那個多的執行紀錄。

小技巧其實很簡單，就是在指令後面加上 | grep "avg loss" ，只顯示有 avg loss 字樣的那些 log 就可以了。所以執行指令會長得像是這樣：

# train the model
!./darknet detector train {LOCAL_CFG_DIR_PATH}/obj.data {LOCAL_CFG_DIR_PATH}/yolov3.cfg {GDRIVE_WEIGHTS_DIR_PATH}/yolov3_last.weights  | grep "avg loss"

完整 Colab 範例程式碼請見 http://bit.ly/2XCOT3E