在自定義的數據集上訓練YOLOv5詳細教程分享

點擊上方「AI算法與圖像處理」，選擇加"星標"或「置頂」

重磅乾貨，第一時間送達

來源 | 深度學習與計算機視覺

編輯 | 磐懟懟

YOLO系列的目標檢測模型隨着YOLOv5的引入變得愈來愈強大了。在這篇文章中，咱們將介紹如何訓練YOLOv5爲你識別本身定製的對象。

本文咱們使用公共血細胞檢測數據集，你能夠本身導出，也能夠在本身自定義數據上使用本教程。

• 公共血細胞檢測數據集： https://public.roboflow.ai/object-detection/bccd

爲了訓練檢測器，咱們採起如下步驟：

• 安裝YOLOv5依賴項
• 下載自定義YOLOv5對象檢測數據
• 定義YOLOv5模型配置和架構
• 訓練一個定製的YOLOv5檢測器
• 評估YOLOv5性能
• 可視化YOLOv5訓練數據
• 對測試圖像使用YOLOv5進行推斷
• 導出並保存YOLOv5權重以供未來使用

YOLOv5的新技術點

就在兩個月前，咱們對googlebrain引入EfficientDet感到很是興奮，並寫了一些關於EfficientDet的博客文章。咱們認爲這個模型可能會超越YOLO家族在實時目標探測領域的突出地位，但事實證實咱們錯了。

三週內，YOLOv4在Darknet框架下發布，咱們還寫了不少關於YOLOv4技術解析的文章。

在寫這些文章的幾個小時以前，YOLOv5發佈了，咱們發現它很是清晰明瞭。

YOLOv5是在Ultralytics-Pythorch框架中編寫的，使用起來很是直觀，推理速度很是快。事實上，咱們和許多人常常將YOLOv3和YOLOv4 Darknet權重轉換爲Ultralytics PyTorch權重，以便使用更輕的庫來更快地進行推理。

YOLOv5比YOLOv4表現更好嗎？咱們很快會向你介紹，在此以前你須要已經對YOLOv5和YOLOv4有了初步的瞭解。

YOLOv5與EfficientDet的性能對比

YOLOv4顯然沒有在YOLOv5存儲庫中進行評估，也就是說，YOLOv5更易於使用，並且它在咱們最初運行的定製數據上表現很是出色。

咱們建議你在 YOLOv5 Colab Notebook 中同時進行接下來的操做。

• https://colab.research.google.com/drive/1gDZ2xcTOgR39tGGs-EZ6i3RTs16wmzZQ

安裝YOLOv5環境

首先咱們克隆YOLOv5存儲庫並安裝依賴項，這會設置咱們的編程環境，準備好運行對象檢測訓練和推理命令。

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
!git clone https://github.com/ultralytics/yolov5 # clone repo!pip install -U -r yolov5/requirements.txt # install dependencies
%cd /content/yolov5

而後，咱們能夠看看谷歌Colab免費提供給咱們的訓練環境。

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
import torchfrom IPython.display import Image # for displaying imagesfrom utils.google_utils import gdrive_download # for downloading models/datasets
print('torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

你會從谷歌Colab收到一個 Tesla P100 GPU。如下是我收到的：

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
torch 1.5.0+cu101 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', major=6, minor=0, total_memory=16280MB, multi_processor_count=56)

GPU可讓咱們加快訓練時間，Colab預裝了 torch 和 cuda 。若是你嘗試在本地使用本教程，可能須要執行其餘步驟來設置YOLOv5。

下載自定義YOLOv5對象檢測數據

在本教程中，咱們將從Roboflow下載YOLOv5格式的自定義對象檢測數據。在本教程中，咱們使用公共血細胞檢測數據集訓練YOLOv5檢測血流中的細胞，你可使用公共血細胞數據集或上傳你本身的數據集。

• Roboflow： https://roboflow.ai/
• 公共血細胞數據集： https://public.roboflow.ai/object-detection/bccd

關於標記工具的說明

若是你有未標記的圖像，則首先須要標記它們。對於免費的開源標籤工具，咱們推薦你閱讀 LabelImg入門 或 CVAT註釋工具入門 的教程指南。嘗試標記約50幅圖像再繼續本教程，由於在之後提升模型的性能的過程當中，你將須要添加更多標籤。

• https://blog.roboflow.ai/getting-started-with-labelimg-for-labeling-object-detection-data/
• https://blog.roboflow.ai/getting-started-with-cvat/

一旦你標記了數據，要將數據移動到Roboflow中，請建立一個免費賬戶，而後你能夠以任何格式拖動數據集：（VOC XML、COCO JSON、TensorFlow對象檢測CSV等）。

上傳後，你能夠選擇預處理和加強步驟：

爲BCCD示例數據集選擇設置

而後，單擊 Generate 和 Download，你將能夠選擇YOLOv5 Pythorch格式了。

選擇「YOLO v5 Pythorch」

當出現提示時，必定要選擇「Show Code Snippet」，這將輸出一個下載curl腳本，這樣你就能夠輕鬆地將數據以正確的格式移植到Colab中。

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
curl -L "https://public.roboflow.ai/ds/YOUR-LINK-HERE" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

正在Colab中下載…

下載YOLOv5格式的自定義對象數據集

導出將會建立一個名爲data.yaml的YOLOv5.yaml文件，指定YOLOv5 images文件夾、YOLOv5 labels文件夾的位置以及自定義類的信息。

定義YOLOv5模型配置和架構

接下來，咱們爲咱們的定製對象檢測器編寫一個模型配置文件。在本教程中，咱們選擇了最小、最快的YOLOv5基本模型，你也能夠從其餘YOLOv5模型中選擇，包括：

• YOLOv5s
• YOLOv5m
• YOLOv5l
• YOLOv5x

你也能夠在此步驟中編輯網絡結構，但通常不須要這樣作。如下是YOLOv5模型配置文件，咱們將其命名爲 custom_yolov5s.yaml ：

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
nc: 3depth_multiple: 0.33width_multiple: 0.50
anchors: - [10,13, 16,30, 33,23]  - [30,61, 62,45, 59,119] - [116,90, 156,198, 373,326]
backbone: [[-1, 1, Focus, [64, 3]], [-1, 1, Conv, [128, 3, 2]], [-1, 3, Bottleneck, [128]], [-1, 1, Conv, [256, 3, 2]], [-1, 9, BottleneckCSP, [256]], [-1, 1, Conv, [512, 3, 2]],  [-1, 9, BottleneckCSP, [512]], [-1, 1, Conv, [1024, 3, 2]], [-1, 1, SPP, [1024, [5, 9, 13]]], [-1, 6, BottleneckCSP, [1024]], ]
head: [[-1, 3, BottleneckCSP, [1024, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], [-2, 1, nn.Upsample, [None, 2, "nearest"]], [[-1, 6], 1, Concat, [1]], [-1, 1, Conv, [512, 1, 1]], [-1, 3, BottleneckCSP, [512, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], [-2, 1, nn.Upsample, [None, 2, "nearest"]], [[-1, 4], 1, Concat, [1]], [-1, 1, Conv, [256, 1, 1]], [-1, 3, BottleneckCSP, [256, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],
 [[], 1, Detect, [nc, anchors]], ]

訓練定製YOLOv5探測器

咱們的 data.yaml 和 custom_yolov5s.yaml 文件已經準備好了，咱們庫開始訓練了！

爲了開始訓練，咱們使用如下選項運行訓練命令：

• img：定義輸入圖像大小
• batch：肯定batch大小
• epochs：定義epochs。（注：一般，3000+很常見！）
• data：設置yaml文件的路徑
• cfg：指定咱們的模型配置
• weights：指定權重的自定義路徑。（注意：你能夠從Ultralytics Google Drive文件夾下載權重）
• name：結果名稱
• nosave：只保存最後的檢查點
• cache：緩存圖像以加快訓練速度

運行訓練命令：

訓練定製的YOLOv5探測器。它訓練得很快！

在訓練期間，你能夠看 mAP@0.5 來了解你的探測器是如何運行的，請參閱這篇文章。

• https://blog.roboflow.ai/what-is-mean-average-precision-object-detection/

評估定製YOLOv5檢測器的性能

如今咱們已經完成了訓練，咱們能夠經過查看驗證指標來評估訓練過程的執行狀況。訓練腳本將刪除tensorboard日誌，咱們將其可視化：

在咱們的自定義數據集上可視化tensorboard結果

若是你由於一些緣由不能把張量可視化，能夠用utils.plot_result來繪製並保存爲result.png。

你須要在驗證評估分數達到其最高點處獲取訓練好的模型權重。

可視化YOLOv5訓練數據

在訓練過程當中，咱們能夠可視化真實訓練數據和加強後的訓練數據。

咱們的真實訓練數據

咱們的訓練數據採用自動YOLOv5加強

對測試圖像運行YOLOv5推斷

如今咱們利用咱們訓練好的模型，對測試圖像進行推理。訓練完成後，模型權重將保存到 weights/。

推理過程，咱們調用這些權重和一個指定模型置信度的conf（要求的置信度越高，預測越少）、以及一個推理源。源能夠接受一個包含圖像、單個圖像、視頻文件以及設備的網絡攝像頭端口的目錄。對於源代碼，我將 test/*jpg 移到 test-infer/ 。

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
!python detect.py --weights weights/last_yolov5s_custom.pt --img 416 --conf 0.4 --source ../test_infer

推理時間很是快，在咱們的 Tesla P100 上，YOLOv5s 達到了每秒142幀！！

以142 FPS（0.007s/圖像）的速度推斷YOLOv5s

最後，咱們在測試圖像上可視化咱們的檢測器推斷結果。

測試圖像的YOLOv5推理

導出並保存YOLOv5權重以供未來推斷

既然咱們定製的YOLOv5物體檢測器已經經過驗證，咱們須要從Colab中取出權重，用於實時計算機視覺任務，爲此咱們導入一個Google驅動器模塊並將其發送出去。

    
    
    
     
     
              
     
     

     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
     
     
     
      
      
               
      
    
    
    
     
     
              
     
     
from google.colab import drivedrive.mount('/content/gdrive')
%cp /content/yolov5/weights/last_yolov5s_custom.pt /content/gdrive/My\ Drive

結論

咱們但願你能夠訓練屬於你本身的定製YOLOv5檢測器！

使用 YOLOv5 很是方便，並且訓練迅速，推理迅速，表現出色。讓咱們把它弄出來！

參考連接：https://blog.roboflow.ai/how-to-train-yolov5-on-a-custom-dataset/

下載1：動手學深度學習

在「AI算法與圖像處理」公衆號後臺回覆：動手學深度學習，即可下載547頁《動手學深度學習》電子書和源碼。該書是面向中文讀者的能運行、可討論的深度學習教科書，它將文字、公式、圖像、代碼和運行結果結合在一塊兒。本書將全面介紹深度學習從模型構造到模型訓練，以及它們在計算機視覺和天然語言處理中的應用。

下載2

在「AI算法與圖像處理」公衆號後臺回覆： OpenCV實戰項目20講 ，便可下載20個有趣的OpenCV實戰項目

    
 
 
     我的微信（若是沒有備註不拉羣！） 

    

    
    
    
     
     
              
     

    
 
 
     請註明： 
 
     地區+學校/企業+研究方向+暱稱 

    

    
    
    
     
     
              
     

    
 
 
     
 

    

本文分享自微信公衆號 - AI算法與圖像處理（AI_study）。
若有侵權，請聯繫 support@oschina.cn 刪除。
本文參與「OSC源創計劃」，歡迎正在閱讀的你也加入，一塊兒分享。