[實作筆記] 在Pytorch框架下以YOLOv4訓練自定義資料集

<source : Kaggle戴口罩的資料集>

0.前言

站在眾多前輩（巨人）的肩膀上看世界。本篇文章主要是分享如何進行YOLOv4模型的訓練資料標記，以及如何在 Pytorch 架構下利用YOLOv4-tiny的模型進行自定義資料集訓練。

1.訓練資料標記

YOLOv4是常見的物件辨識模型之一，為了讓模型能有效學習，必須對訓練資料進行標記，也就是告訴模型，模型需要辨識的物件是什麼，以及該物件位於圖片中的哪個位置。

本文將以Labelimg進行說明。其步驟如下：

1.1.安裝與啟動圖片標記工具Labelimg

在指定虛擬環境中輸入指令 pip install labelimg 。若您尚未建立指定虛擬環境，可參考< Pytorch (Windows) 環境建置 與 YOLOv4 起手式實作 >。

安裝完畢後，在指定虛擬環境中輸入指令 labelimg ，啟動Labelimg。如下圖所示。

1.2.進行圖片標記

啟動Labelimg後，操作介面如下所示：

source : Kaggle戴口罩的資料集

接著針對圖片中您想標記的物件，用滑鼠進行框選，並給予為該類別命名。本文以戴口罩的人為例，類別命名為 "mask"。如下圖所示。所有需標記的圖片比照此模式進行。

source : Kaggle戴口罩的資料集

每張圖片標記後會存成XML檔。

2.在Pytorch框架下以YOLOv4訓練自定義資料集

站在巨人的肩膀上看世界，本例使用此作者 bubbliiiing 的Github 進行自定義資料訓練。步驟如下：

2.1. 下載YOLOv4範例程式與預訓練模型權重檔

請在指定虛擬環境中輸入指令 git clone https://github.com/bubbliiiing/yolov4-tiny-pytorch.git

請至原作者github下載pretrained model ( yolov4_tiny_weights_coco.pth )，放在model_data的資料夾中。

source：Greatway9999

2.2. 自定義資料處理

本例YOLOv4模型指接受的訓練資料為jpg的圖片格式，因此輸入圖片如果尾碼非jpg，需要自己批量轉成jpg後再開始訓練。

假設您已經把自定義的資料都標註好了。本例使用的資料集為Kaggle戴口罩的資料集，如果還沒準備好資料，可至Kaggle下載。該資料集總共有852張圖片包含XML標記檔，可節省您準備資料的時間，幫助您有多餘的心力來學習如何實現用YOLOv4進行自定義資料集訓練。

本文使用VOC格式進行訓練，請將準備好的資料將標記好的文件放在VOCdevkit資料夾中。訓練前將圖片檔放在VOCdevkit資料夾下的VOC2007資料夾下的JPEGImages中。

source：Greatway9999

訓練前將標籤文件放在VOCdevkit資料夾下的VOC2007資料夾下的Annotation中。

source：Greatway9999

2.3.訓練前的參數調整

由於是自定義資料集，因此類別等資訊都和原作者所提供的內容不同，所以要進行相關參數調整。

在完成資料集的擺放之後，我們需要對資料集進行下一步的處理，目的是獲得訓練用的2007_train.txt以及2007_val.txt，需要用到根目錄下的 voc_annotation.py。

voc_annotation.py裡面有一些參數需要設置。分別是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path，第一次訓練可以僅修改classes_path。

在model_data資料夾中，創建 cls_classes.txt 的類別資料，內容為自定義資料夾類別，本例為 "with_mask" 和 "without_mask"。

source：Greatway9999

接著以下3份文件檔均需設定 cls_classes.txt 之路徑：

• voc_annotation.py (如下圖所示)
• train.py
• yolo.py

執行 voc_annotation.py，生成了2007_train.txt以及2007_val.txt。此時可以開始訓練了。輸入指令　python voc_annotation.py

source：Greatway9999

2.4. 訓練

可依據需求調整 train.py 中的Optimizer、Epoch。此例維持原作者預設參數。

確認無誤後，輸入指令 python train.py

接著就是等待模型訓練完畢。以筆者的電腦環境，訓練時間約3個多小時。

2.5. 結果預測

訓練完的權重檔會存放在 log/ 資料夾中，請開啟 yolo.py ，並在"model_path"處更新為最新的權重檔的路徑。

接著輸入指令 python predict.py ，即可進行測試。

測試結果如下圖所示。

<source : Kaggle戴口罩的資料集>

-本文結束-

---

參考資料

• LabelImg 影像標註工具使用教學，製作深度學習用的資料集
• 如何轉換為Yolo txt格式
• Github:bubbliiiing/yolov4-tiny-pytorch
• Github:bubbliiiing/yolov4-pytorch
• PyTorch Yolov4配置
• 深度學習-物件偵測YOLO cfg檔解讀 (三) - 2021年
• bubbliiiing/yolov5-v6.1-pytorch
• 睿智的目标检测30——Pytorch搭建YoloV4目标检测平台
• [機器學習 ML NOTE]YOLO!!!如何簡單使用YOLO訓練出自己的物件偵測!!! (Windows+Anaconda)
• 【YOLOv4-pytorch】训练自己的数据集实践记录及问题总结
• Kaggle戴口罩的資料集