天天看點

從零開始帶你一步一步使用YOLOv3訓練自己的資料

YOLOv3是比較常見和常用的深度學習目标檢測(Object Dection)算法。今天給大家介紹一下如何一步一步使用YOLOv3訓練自己的資料集。

從零開始帶你一步一步使用YOLOv3訓練自己的資料

一、标注資料集

首先我們需要使用 labelimg 工具來标注圖檔資料集,例如圖檔是 .jpg 格式的,用矩形框标注圖檔中的目标位置,得到 .xml 檔案。這裡 labelimg 的使用方法就不作介紹了,讀者自行查閱。

從零開始帶你一步一步使用YOLOv3訓練自己的資料

例如 100 張圖檔對應得到 100 個 .xml 檔案:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

得到 .jpg 和 .xml 檔案之後,我們還需要把資料集整理成 VOC2007 規定的格式,友善我們使用 YOLOv3 進行訓練。

方法很簡單,建立如下目錄層級:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

然後将 100 張 .jpg 圖檔放入到 VOCdevkit/VOC2007/JPEGImages 目錄下,将 100 個 .xml 檔案放入到 VOCdevkit/VOC2007/Annotations 目錄下。

至此,資料集準備完畢!

二、下載下傳 YOLOv3 源碼

在我們的 Ubuntu 系統或者伺服器上,使用 git 指令直接下載下傳 YOLOv3 的源碼工程:

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

下載下傳完成之後,進入 darknet 目錄,會發現如下:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

YOLOv3 使用的是開源的神經網絡架構 Darknet53,有 CPU 和 GPU 兩種模式。預設使用的是 CPU 模式,如果我們使用 GPU 的話,需要修改 darknet 目錄下的 Makefile 檔案:

$ vim Makefile      

修改的幾點如下:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

修改的地方有 4 處。其中,第 1 處 GPU=1,CUDNN=1,表示使用 GPU 和 CUDNN,若使用 CPU 的話,置為 0 即可;第 2 處将 NVCC 設為本地安裝 nvcc 的實際目錄;第 3 處将 COMMON+ 設為本地 cuda 的頭檔案目錄;第 4 處将 LDFLAGS+ 設為 cuda 的庫目錄。

Makeflie 檔案修改完成之後,儲存退出。

三、準備資料

将我們之前準備好的包含 100 張 .jpg 圖檔和 100 個 .xml 檔案的 VOCdevkit 檔案夾拷貝到 darknet 目錄下,替換原來的 VOCdevkit 檔案夾。

進入到 darknet/VOCdevkit/VOC2007/ 目錄下,建立 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()      

該 test.py 腳本的作用就是對 100 個樣本進行訓練集和測試集的劃分,trainval_percent = 0.8 表示劃分 80% 的樣本作為訓練集,20% 的樣本作為測試集。

在目前目錄,運作該腳本:

$ python test.py      

運作完之後,在 /ImageSets/Main/ 目錄下會生成 4 個檔案:test.txt、train.txt、trainval.txt、val.txt。其中,trainval.txt 存放的是 80 個訓練集的樣本名,test.txt 存放的是 20 個測試集的樣本名。

從零開始帶你一步一步使用YOLOv3訓練自己的資料

傳回到 drknet 目錄下,建立 voc_label.py 腳本,也可以使用下面指令下載下傳:

$ wget https://pjreddie.com/media/files/voc_label.py      

voc_label.py 腳本内容如下:

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["person", "car", "bicycle", "train"]


def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(year, image_id):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()

for year, image_set in sets:
    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()      

該檔案需要修改兩個地方,分别是 sets 和 classes,classes 根據實際的類别進行修改。

修改完成之後,儲存退出。在 darknet 目錄下,運作該腳本:

$ python voc_label.py      

運作之後,在目前目錄下生成 3 個檔案:2007_train.txt, 2007_val.txt, 2007_test.txt,存放的是訓練集、測試集圖檔的路徑。

運作下面的指令:

$ cat 2007_train.txt 2007_val.txt  > train.txt      

這樣,train.txt 即為真正的訓練集圖檔路徑,2007_test.txt 即為真正的測試集圖檔路徑。

同時,在 VOCdevkit/VOC2007 目錄下也會多生成一個 labels 檔案夾,labels 檔案夾裡存放的就是訓練集每個圖檔包含的類别、矩形框四個坐标資訊。這時候資料集正式完成。

四、修改配置檔案

1. 修改 voc.data 檔案

在 drknet 目錄下,打開 cfg/voc.data 檔案,根據實際情況作如下修改:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

其中,class 是類别,train 和 valid 是之前運作 voc_label.py 得到的 train.txt 和 2007_test.txt。names 改為自己的實際路徑,backup 為模型存放的路徑。

2. 修改 voc.names 和 coco.names

在 darknet 目錄下,打開 data/voc.names 和 data/coco.names 檔案,修改成自己的類别:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

儲存退出!

3. 修改 yolov3-voc.cfg

下面需要修改配置檔案,打開 cfg/yolov3-voc.cfg,搜尋 yolo, 總共會搜出 3 個含有 yolo 的地方。

每個地方都必須要改 2 處:

  • 參數 filters 由下式計算:3*(5+classes),例如本例中 classs=4,則filters=27;
  • 參數 class 改為實際的類别個數;
從零開始帶你一步一步使用YOLOv3訓練自己的資料

yolov3-voc.cfg 檔案中修改 filters 和 classes 的地方一共有 3 處,注意别遺漏了。

在 yolov3-voc.cfg 檔案的開頭處,修改訓練的一些參數:

從零開始帶你一步一步使用YOLOv3訓練自己的資料

一般根據具體情況作适當修改即可,注意,訓練的時候,Testing 的 batch 和 subdivisions 需要注釋掉,learning_rate 是學習率,max_batches 是最大疊代訓練次數,可根據訓練集大小自行修改。

關于 yolov3-voc.cfg 中更詳細的參數解釋,可自行查閱。

五、編譯

剛才我們修改了 Makefile檔案,并修改了各個配置檔案,下面就對該工程進行編譯,在 darknet 目錄下輸入下面的指令:

$ make      

編譯完成之後,生成 darknet 可執行程式。

六、訓練

在 darknet 目錄下,

首先下載下傳 YOLOv3 的預訓練模型:

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

然後直接輸入下面的指令進行訓練:

$ ./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg darknet53.conv.74 -gpus 0,1      

其中,-gpus 0,1 表示使用 2 塊 gpu 進行訓練,編号是 0 和 1。根據你可使用的 gpu 個數進行調整。

從零開始帶你一步一步使用YOLOv3訓練自己的資料

當看到上面的列印内容,基本就表示訓練開始了,慢慢等待吧~

關于 YOLOv3 詳細的訓練步驟就介紹到這裡,下一篇我将帶大家來學習一下使用 YOLOv3 訓練好的模型進行實際的圖檔測試還有如何進行批量測試~

繼續閱讀