Yolov5訓練自有模型


之前一篇關於yolo的博客快速的將yolov5的demo運行起來體驗了一下深度學習的威力,這篇博客以自己搭建口罩識別為例,主要記錄一下嘗試yolov5訓練自有模型並成功運行的過程。

我一開始准備在windows上測試的,但那台筆記本空間內存都有限,無法實現訓練,只能測試訓練好的模型。后來選擇在自己的Mac Pro筆記本上面進行測試。要說明的一點是,在Mac上同樣要准備好Tensorflow、PyTorch、OpenCV(版本要注意),標注軟件LabelImg還需要PyQt等環境,但都可以通過百度解決。另外就是推薦使用Anaconda安裝,會比較方便的管理python的環境。

一、數據標注

這里需要標注軟件LabelImg(PyQT開發,支持多系統)。

將圖片分門別類存放,標記好之后會同步生成圖片文件和txt標記文件。

  

生成的txt文件里的數據:

PS:標注是一個體力活,需要耐心(需要標注的圖片和數量很多)和細心(標注時要盡量准確)。圖片的數量及標記的質量是識別模型的基礎和關鍵。

二、訓練模型

這里極度推薦官方文檔的說明:https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

過程非常簡潔明了,建議初學者一定要細心閱讀這份文檔,按照步驟完成模型的訓練。

模型訓練的命令格式如下:

python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt

train.py是yolov5根目錄下的一個py文件,所以需要先切換到yolov5根目錄下才能執行上面的命令。

上面還引用了一個coco128的yaml文件,內容如下。

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
nc: 80  # number of classes
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']  # class names


# Download script/URL (optional)
# download: https://ultralytics.com/assets/coco128.zip

這是一個示范文件,按照這個模板格式,我們在訓練自己的模型時也需要創建這個yaml文件。自己修改里面設置的路徑和類別等信息。其它參數設置相信都能看得懂。

當然這里是采用的cpu訓練,訓練的速度跟圖片的數量和大小都有關系,如果數據量大的話訓練幾天的都有。

這里預留兩個問題以后驗證:

1. 就是中斷訓練過程后能否接着繼續訓練。

2. 嘗試gpu訓練看速度是否很快。

訓練完成之后會在yolov5根目錄下的runs/train下生成對應的exp目錄。

可以看到訓練的過程。里面的weights是生成的權重文件,分別是best.pt和last.pt。我們拷貝best.pt到yolov5的根目錄。

在終端切換到yolov5根目錄,執行如下命令即可驗證視頻流檢測我們的自訓練模型的效果啦~

python detect.py --source 0 --weights best.pt --conf 0.4

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM