yolov5實戰之皮卡丘檢測

haoliuhust發表於2021-01-14
YOLO

前言

從接觸深度學習開始一直都做的是人臉識別，基本上也一直都在用mxnet. 記得之前在剛接觸的時候看到部落格中寫到，深度學習分三個層次，第一個層次是分類，第二個層次是檢測，第三個層次是分割。人臉識別算是分類問題，也就是一直在第一個層次···一直都想有機會了解下第二個層次，奈何精力有限，工作中也沒有實際的專案需要。最近正好有個不急的檢測專案，趁此機會入門檢測吧。工作中聽同事提到yolov5效果不錯，而且文件指導也比較豐富，因此選擇從此入手，順便也熟悉下pytorch。本文就以訓練一個簡單的皮卡丘檢測作為入門吧，也激發下對檢測的學習興趣，暫時不涉及網路，anchor等細節的理解，只以訓練完為目標。

環境準備

安裝torch這些就不說了，本次訓練基於yolov5官方專案：
https://github.com/ultralytics/yolov5
clone後，pip install -r requirements.txt安裝各項依賴
基本的目錄結構如下：
data: 訓練資料和資料配置檔案以及訓練配置
models: 模型配置檔案
utils: 資料讀取、nms等各種訓練測試中用到的指令碼
weight:
train.py: 訓練指令碼
test.py：測試指令碼
detect.py: 執行指令碼

資料集準備

https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
根據官方指導，按需要的格式準備資料集：

  1. 一個圖片檔案對應一個同名的.txt label檔案
  2. label檔案中每一行表示一個目標，格式為：class x_center y_center box_width box_height, 其中數值都是歸一化的。類別從0開始
  3. 按類似目錄方式防止資料，yolov5通過將圖片路徑中最後一個images換成labels去尋找對應的標籤
coco/images/000000109622.jpg  # image
coco/labels/000000109622.txt  # label

本文中用的資料來源於 目標檢測資料集（皮卡丘） 中，但原文是mxnet用的rec格式，因此需要轉換成yolo格式。我已經轉好了：
連結: 百度網盤資料集 提取碼: rek1
下載後解壓即可得到符合格式要求的資料。
資料示例：

image

訓練前的配置

修改或新建dataset.yaml

在data目錄下，我們可以看到一個coco.yaml，內容如下：

# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /coco128
#     /yolov5


# download command/URL (optional)
# download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./data/coco128/images/train2017/  # 訓練集圖片路徑
val: ./data/coco128/images/train2017/  # 驗證集圖片路徑

# number of classes，類別數
nc: 80

# class names，類別名
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

仿照這個格式，建立皮卡丘資料集的配置檔案：

# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /coco128
#     /yolov5


# download command/URL (optional)
# download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./data/pikachu/images/train/  # 128 images
val: ./data/pikachu/images/val/  # 128 images

# number of classes
nc: 1

# class names
names: ['pikachu']

訓練引數配置

訓練引數配置在data/hyp.finetune.yaml和hyp.scratch.yaml中，前者是finetune用的，後者是從頭訓練的。這次我們直接用hyp.finetune.yaml，先不做修改。

模型配置

模型配置在models/yolov5s.yaml中。目前也不做修改。

開始訓練

訓練使用train.py指令碼，主要引數有：

--weights 預訓練權重，“”表示從頭訓  
--cfg 模型配置檔案路徑  
--data 資料配置檔案路徑  
--hyp 訓練引數配置檔案路徑  
--epochs 訓練的epoch數 
--batch-size 
--img-size 網路的輸入大小  
--noautoanchor 是否自動計算anchor  
--device 裝置號  
--workers dataloader執行緒數

還有很多其他引數，不過暫時用不上就先不管了,執行以下命令可開始訓練，指令碼會自動去下載yolov5s的預訓練權重進行訓練。

python train.py --weights yolov5.pt --data data/pikachu.yaml --hyp data/hyp.finetune.yaml --epochs 40 --img-size 320 --device 0 --batch-size 512

視覺化

yolov5裡面有兩種視覺化方式，一種是tensorboard, 一種是wandb。
其中wandb有兩種使用方式，一種是到https://wandb.ai/home註冊一個賬號，線上使用。註冊後新建專案,名字叫yolov5，然後本地安裝配置wandb, 按提示輸入必要的資訊（API Key):

pip install wandb  
wandb login

那麼訓練時就可以在自己的專案中看到如下畫面，還是挺不錯的，功能挺豐富：
image
不過這個網站挺卡的，好在wandb也有本地使用方式。

安裝docker後
wandb local

參考：https://docs.wandb.ai/self-hosted/local
然後就可以本地訪問了。

測試

有兩個指令碼和測試有關，test.py是用來跑測試集測效能的，detect.py是用來測試資料夾或攝像頭的。

python detect.py --source data/images --weights yolov5s.pt --conf-thres 0.25

在run/detect下即可找到影像畫有目標框的影像（也可以在執行detect.py的時候指定--view-img選項，即可顯示影像）：
image

結語

至此就訓練出了一個簡單的檢測模型呢，雖然很簡單，不過基本流程通了。在走通流程後就要繼續瞭解實現的細節，如何調參等等了。可以從瞭解網路結構開始，可參考 [https://zhuanlan.zhihu.com/p/343195876](https://zhuanlan.zhihu.com/p/343195876

注：如果資料集連結過期了，關注 老司機的視覺屋，回覆pikachu即可獲取連結

相關文章