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