yoloV5-6.2分类训练记录

加我微信拉你进群交流：wu331376411

一 背景介绍

yoloV6，V7相继跟新，没有想到用的最熟悉的V5又双叒叕更新了，今天我就来给大家准备分享一下yoloV5-6.2的分类训练。

二 模型下载

首先从官网下载yoloV5的最新代码。
git地址：https://github.com/ultralytics/yolov5

下载后的代码结构如下：

我们可以看到新增加了classify（分类）和segment（分割）两部分的训练代码，今天我们就来看看分类代码
这里我们还需要下载一下官方的分类与训练模型：
yolov5s-cls.pt: https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s-cls.pt

三 训练记录

1 添加与训练模型：

在yolov5/classify/train.py中，将下载好的预训练模型位置添加一下即可。

2 分类数据整理：

这里不需要做txt， 这里不需要做txt ，这里不需要做txt ，重要的事情说三遍。
直接把分类的路径记录一下即可：

这里只需要有test或者val 和train就可以了，也可以做三个文件夹，都行。至少2个文件夹。
每个文件夹里面是对应的类别文件夹。

4个类别就写四个文件夹，一个文件夹存放一个类别的数据，train和test是一样的准备就可以了。

注意： 大家关心的问题，给大家也实验了一下，如果一个类别的数据非常多，导致文件夹里面存放会卡死。
比如：red文件夹里面有10万的数据，能不能分文件夹，答案是：可以。

可以按照不同的文件夹分类存放数据，一样可以读取到所有的数据，并且都属于red这一个类别。

整理好的数据，将总根目录写到配置文件处就可以了：

这里是写到yolo_test，也就是train，test的上级目录。

准备好这些内容调整epoch和batch-size就可以开始训练了。是不是很简单。

四 注意事项

1 报错一：AttributeError: ‘numpy.ndarray’ object has no attribute ‘imshow’

这个报错是由于utils/plots.py 中代码有点问题。总的来说也是自己batch-size写的太小导致的。

• 修改方法一： 增加batch-size 我最开始设置的为2 ，这里改到4或者8 就好了，如果8还不行，就继续增加。直到不报错。

• 修改方法二：
  这里报错主要是batch-size太小，导致m是等于1 的，无法对ax做ravel()操作报错。所以在代码 utils/plots.py 457行，添加一个等于号就解决问题。
  

2 训练图片resize不正确：

开始训练的过程，发现图片都是随机裁剪了的，有可能导致训练数据了中重要的信息丢失。这样会导致分类模型后期的分类出差。也是非常致命的问题。

这里作者在数据导入部分分别使用的是随机裁剪和中心裁剪的方式，所以我们需要改一下。
在utils/dataloaders.py 中，主要是这2个transforms需要修改一下。

大家对照 utils/augmentations.py 代码里面行号我框出来的位置，修改一下就可以了，分别是CenterCrop和classify_albumentations，2个函数里面的裁剪方法。

修改后训练截图：图像有些变形，但是显示的是全部信息。这样就没有什么问题了。

写在最后

• 主要是图像resize的问题，会严重影响训练结果的预测，导致出现预测不准的情况。
• 做分类只需要把数据存放到同类别的文件夹下，省去了做txt的部分，确实也简化了一些步骤和工作量。