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

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章