YOLO-V5分类实战系列 —— 快速训练自己的数据集

---

1、获取官方源码

官方代码：https://github.com/ultralytics/yolov5
更新版本：v6.2 添加了分类模块，v7.0 添加分割模块，所以至少要下载 v6.2 之后的版本
环境安装：按照下面的方式安装，（本人通常使用conda软件，首先用conda安装pytorch，然后使用pip安装其它依赖库）
 
拉取特定版本的代码：参考我的博文 YOLO-V5 算法和代码解析系列（一）—— 快速开始

2、测试官方源码

下图是官方的代码结构（v7.0），文件夹【classify】存放分类的相关代码，包括训练代码 train.py，预测代码 predict.py，评估代码 val.py，

---

2.1、公开数据集测试源码

按照上述代码结构，进入【classify】文件夹，打开训练脚本【train.py】。根据训练脚本最上面的注释内容：可以使用官方数据集，也可以使用自己的数据集，以及单卡和多卡的训练命令，

官方训练命令如下，

python classify/train.py --model yolov5s-cls.pt --data cifar10 --epochs 5 --img 224

运行上述命令，会【自动下载】所需的数据集【cifar10】和模型文件【yolov5s-cls.pt】。下载完成后，数据的存放结构如下图所示，包含测试集【test】和训练集【train】。注意，参考官方的数据存放结构，后续存放自己的数据集

如下图所示，每一个文件夹内，存放一个类别的图片

训练过程如下，训练基本正常，

---

3、源码模块解析

3.1、数据读取

• 训练类别数量：数据路径【data_dir/train】下的文件夹数量等同于训练的类别数量（存放数据时应该注意），具体代码如下：

  nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()])  # number of classes
  

• 数据读取

  需要注意的几个重要参数
  path : 数据存放位置，数据放置完成，该参数不需要配置，默认即可
  imgsz：输入到网络训练的图片尺寸
  batch_size：批次数量
  augment：是否进行数据增强，通常训练阶段为【true】，测试阶段为【false】

  

  数据读取的核心代码，包括数据读取和数据增强。

  

  数据增强：self.album_transorms【第三方数据增强库】，self.torch_transforms【torch自带库】；如果添加数据增强操作，需要在这两个函数中添加，

  
  

---

3.2、网络结构

分类网络【yolov5-cls 】的网络结构如下图所示，只用了目标检测网络的主干部分。基础模块包括：C3，CBS，其详细的结构，可以参考之前写的博文，YOLO-V5 系列算法和代码解析（四）—— 网络结构

---

4、快速开始训练自己的数据

按照在官方数据集上的测试过程，要想快速开始训练，只要两方面的准备：（1）准备好自己的数据集；（2）配置部分训练参数。具体的过程在以下章节陈述

4.1、准备自己的数据

根据【2.1】节的陈述，依据官方数据的组织结构来存放自己的数据集，每一个文件夹放一个类别的图片。【mono_demo】就是存放自己的训练数据的根目录名称，也是训练时传入的数据标志。存放的结构如下图所示，

---

4.2、配置训练参数

为了训练自己的数据集，首先保证官方代码能够在自己的数据集上正常训练，并且达到基本的训练效果即可（判断标志就是网络能够正常收敛），之后再进行网络调优。所以，这一章节的目的是让官方源码在自己的数据集上快速开始训练，并且能够得到基本的训练效果。为此，修改训练脚本（train.py）的部分训练参数，具体如下，

• 模型选择，默认是 yolov5-cls.pt，也可以是 efficient-b0,b1,b2,b3，resnet18，

  parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
  

• 修改训练数据路径

  parser.add_argument('--data', 
                      type=str, 
                      default='mono_demo', 
                      help='cifar10, cifar100, mnist, imagenet, etc.')
  

• 根据自己的需要，修改如下参数

  parser.add_argument('--epochs', type=int, default=10)
  parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
  parser.add_argument('--imgsz', '--img', '--img-size', 
  					type=int, 
  					default=224, 
  					help='train, val image size (pixels)')
  

• 根据是否需要预训练，修改参数

  parser.add_argument('--pretrained', nargs='?', const=True, 
  					default=True, 
  					help='start from i.e. --pretrained False')
  

配置完参数后，开启训练，过程如下，从精度来看，基本正常，达到最基本的效果，

---

4.3、模型推理和评估

训练完成模型后，通常要经常进行推理和评估，根据结果评估模型，并且作为后续优化的参考

模型预测：

python classify/predict.py --weights runs/train-cls/exp24/weights/best.pt --source im.jpg

模型评估：

python classify/val.py --weights runs/train-cls/exp24/weights/best.pt --data datasets/biaozhu_train_1

训练可视化：

tensorboard --logdir=runs/train-cls/exp34