安裝與測試

最近YOLOv5最新更新升級到v6.x版本，工程簡便性有提升了一大步，本教程教你基于YOLOv5框架如何訓練一個自定義對象檢測模型，首先需要下載對應版本：

https://github.com/ultralytics/yolov5/releases/tag/v6.1

鼠標滾到最下面下載源碼zip包：

https://github.com/ultralytics/yolov5/archive/refs/tags/v6.1.zip

下載完成之后解壓縮到：

D:pythonyolov5-6.1

首先運行下面命令行完成依賴庫的安裝：

等待一段時間，完成全部安裝之后，運行測試：

python detect.py --weights yolov5s.pt --source dataimageszidane.jpg

運行結果如下：

數據集準備與制作

自己百度收集了一個無人機與飛鳥的數據集，其中訓練集270張圖像，測試集26張圖像。

使用labelImg工具完成標注，工具下載地址：

https://gitee.com/opencv_ai/opencv_tutorial_data/tree/master/tools

YOLOv5要求的數據集目錄結構如下：

其中images是所有圖像的集合，labels是所有標簽信息集合。train表示訓練集、valid表示測試集。這里需要注意的是，labels信息中的標簽信息YOLO格式標注框需要把原始標注信息的標注框格式從：

Left top right bottom

轉換為

Center_x, center_y, width, height

并歸一化到0~1之間，這部分我寫了一個腳本來完成label標簽的生成，把xml的標注信息轉換為YOLOv5的labels文件，這樣就完成了數據集制作。最后需要創建一個dataset.ymal文件，放在與data文件夾同一層，它的內容如下：

# train and val datasets (image directory or *.txt file with image paths)train: uav_bird_training/data/images/train/val: uav_bird_training/data/images/valid/
# number of classesnc: 2
# class namesnames: ['bird', 'drone']

三：模型訓練

制作好數據集之后，模型訓練就成為一件很簡單事情，一條命令行搞定。運行下面的命令行：

python train.py --img 640 --batch 4 --epochs 25 --data uav_bird_trainingdataset.yaml --weights yolov5s.pt

其中uav_bird_training文件夾里是制作好的數據集。這樣就開始訓練，訓練過程中可以通過tensorboard來查看可視化的結果，

PR曲線說明訓練效果還錯！

不同框架與硬件平臺推理比較

YOLOv5的6.x版本支持不同框架模型導出與推理，看下圖：

這里分別導出模型為ONNX、XML、engien格式之后，使用同一段視頻在我的筆記本上（CPUi7, GPU3050ti）：在下面四種推理框架上運行

-OpenCV DNN-OpenVINO-ONNXRUNTIME-TensorRT

測試，分別截圖如下：

OpenCV DNN推理速度

OpenVINO平臺上的推理速度

ONNXRUNTIME GPU推理速度

TensorRT框架部署-FP32版本模型推理統計：

YOLOv5的6.x版本，是支持TensorRT 量化到FP16模型直接導出的，但是不支持INT8量化生成，所以自己實現了導出量化INT8版本，測試結果如下：

python detect.py --weights uav_bird_training/uav_bird_int8.engine --data uav_bird_training/dataset.yaml --source D:/bird/bird_fly.mp4

TensorRT框架部署-INT8版本模型推理統計：

最終比較：

注意：以上結果是基于Python語言上測試結果（前后處理都使用pytorch完成）

番外篇：C++推理與比較

使用C++部署時候，前后處理都改成了基于OpenCV 完成，使用CPU完成前后處理，OpenVINO+CPU運行速度截圖如下：

OpenCV DNN + CUDA版本推理

TensorRT-FP32模型推理速度

TensorRT-INT8模型推理速度

最后把自定義訓練導出的模型轉換為TensorRT Engine文件之后，部署到了我的一塊Jetson Nano卡上面，實現了邊緣端的部署，有圖有真相：

審核編輯 ：李倩