YOLOv5模型训练

一、准备

要求：Python>=3.8以及PyTorch>=1.7

$ git clone https://github.com/ultralytics/yolov5  # clone repo
$ cd yolov5
$ pip install -r requirements.txt  # install

二、创建dateset.yaml

这里我们使用COCO128训练集，该训练集有128张图片用于训练和验证。
数据集配置文件定义如下：

# 训练集和验证集目录
train: ../coco128/images/train2017/
val: ../coco128/images/train2017/

# 分类数
nc: 80

# 分类名称
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']

三、创建标签

可以通过CVAT，makesense.ai，Labelbox平台对图片做先行打标，之后导入为YOLO格式，每一张图片对应一个*.txt标签文件（如果该图片没有目标则不需要*txt文件）。

图片示例：

以上图片标签示例：

四、组织目录

以下为我的目录截图，如果要修改训练集和标签的目录需要同时修改dataset.yaml文件相应配置。

五、选择模型

开始训练之前需要选择一个预训练模型，在这里我选择YOLOv5s，这是一个更轻量级以及快速的模型。
以下列举了各模型的数据：

六、训练

在训练之前需要下载预训练权值文件yolov5s.pt。
训练命令如下：

# Train YOLOv5s on COCO128 for 5 epochs
$ python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov5s.pt

所有训练结果保存在/runs/train目录下。

七、视觉化

安装wandb

$ pip install wandb
$ wandb login  #登录之后再使用第六节训练命令训练一遍即可

注意：登录wandb之前需要在wandb平台注册账号