在yolov5上训练自定义数据集

该文章基于github开源的 ultralytics/yolov5 pytorch实现。本文只适用于Windows。

环境

• OS: Windows10 x64
• 内存: 16G
• CPU: i7-6820HQ
• 软件环境: Conda、python 3.8

步骤

1、环境准备

下载安装conda，博主使用的是MiniConda；在其中创建一个新的环境，环境中的python版本不得低于3.8，此为强制要求，若不满足则可能出现未知问题。这里不再展开介绍，操作步骤请百度或我后期会出新的文章介绍。

2、克隆GitHub仓库

首先确保本地安装有git软件。打开gitbash，输入以下命令

git clone https://github.com/ultralytics/yolov5.git

稍等一会仓库就会克隆下来，若出现网络问题可以考虑多尝试几次或科学上网，这里不再展开。

3、下载数据集

这里使用的数据集是 coco128 , 这个数据集是 coco 数据集的子集，里面带有label。请到 这里 下载该数据集，下载该数据集需要kaggle账号，甚至可能需要科学上网，请自行解决。下载完成并解压。在刚才克隆的yolo仓库的同级目录下面，新建一个文件夹，命名为 Dataset ，将解压后的数据集放入其中，如下图所示：

4、创建数据文件

在yolo的根目录中创建一个文件，命名为 coco128.yaml ，打开该文件，将以下内容复制进去

# 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: ../Dataset/coco128  # dataset root dir
train: ../Dataset/coco128/images/train2017  # train images (relative to 'path') 128 images
val: ../Dataset/coco128/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

其中，path是数据集的根目录，这里使用相对路径；train是用于训练的数据；val是用于验证的数据；test是用于测试的数据；nc是类别数量；name是类别名称。请根据自己的数据集情况自行修改。

5、安装yolov5

打开conda，切换到要安装yolo的环境以及yolo根目录，运行以下命令

pip install -r requirements.txt

安装所有关于yolo的依赖环境

6、下载权重文件（模型）

预训练模型 在此，请选择一个进行下载，这里选择最小的 YOLOv5s ，下载完成之后，会得到一个 yolov5s.pt 文件，将其放置在yolo的根目录下面即可。

7、开始训练

打开conda，切换到要安装yolo的环境以及yolo根目录，运行以下命令

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

这里指定使用 coco128.yaml 中的数据集，并指定使用yolov5s这个模型，后面均可以进行修改。出现下图请选择3，暂不输出

这是一个训练过程可视化的工具，可以实时看到训练过程，这样就能在训练时不无聊，当然，也可以使用 tensorboard。接下来还会输出网络结构，如下图

以及训练过程

我这里不小心断电中断了一下，可以继续训练。

8、训练完成

训练结果全部保存在 runs 这个文件夹下面，可以自行查看。