请先看这篇文章：本博打开方式！！！请详读！！！请详读！！！请详读！！！_Cat-CSDN博客

很多刚入门的粉丝私信我说，他们的电脑配置很好，想在本地训练YOLOv5的自定义数据集，这样方便很多，不用担心资源被收回的问题，那就安排一下。

因为使用了Pycharm编辑器进行训练并对官方代码进行了细微修改，所以请看这篇文章：在Windows 10系统下，如何在Pycharm中配置YOLOv5虚拟环境？_Cat-CSDN博客_pycharm安装yolo

本文章的大部分内容与这篇Yolov5如何在Colab中训练自定义数据集？（更新时间：.3.14）_Cat-CSDN博客相似，所以先阅读上面这篇文章，到了“三、训练”的“4、正式训练”时，请阅读本文。

一、数据集目录结构

将构建好的数据集复制到与yolov5-master同一级目录下，如下图所示。

二、修改yaml文件夹内容

这里一共有2个yaml文件需要修改，一个在data文件夹里，另一个在models文件夹里。

刚才我们将自己的数据集文件名改为coco128，那么此时就要到data里修改coco128.yaml文件内容了，具体的目录结构、修改内容如下图所示：

我们训练时使用的预训练权重如果是yolov5s.pt，那么就得到models目录下的yolov5s.yaml文件中修改相应的参数，其它同理。比如使用的是yolov5s.pt，那么就得去修改，如下图所示。注意，请将预训练模型提前下载好，并放入weights文件夹中。

三、修改train.py代码让训练顺利进行

在Pycharm中打开train.py文件后，找到主函数入口，需要对里面的参数进行修改。下面直接贴出大部分参数所代表的含义。

""" --weights：预训练权重文件的路径--cfg：模型yaml文件的路径--data： 数据yaml文件的路径--hyp：超参数文件的路径--epochs： 训练轮次--batch-size： 每批次文件投入量的多少--img-size：图片尺寸--rect:是否采用矩形训练，默认False--resume: 接着上次被打断训练的结果继续训练--nosave: 不保存模型，默认False--notest: 不进行test，默认False--noautoanchor: 不自动调整anchor，默认False--evolve: 是否进行超参数进化，默认False--bucket: 谷歌云盘bucket，一般不会用到--cache-images: 是否提前缓存图片到内存，以加快训练速度，默认False--image-weights：使用加权图像选择进行训练--device: 训练的设备，有CPU和GPU，0代表一个设备--multi-scale: 是否进行多尺度训练，默认False--single-cls: 数据集是否只有一个类别，默认False--adam:是否使用adam优化器--sync-bn: 是否使用跨卡同步BN,在DDP模式使用--local_rank： DDP参数，请勿修改--workers：最大工作核心数--project: 训练模型的保存位置--name： 模型保存的目录名称--exist-ok：模型目录是否存在，不存在就创建"""

上述的训练参数都有默认值，但是每台训练的设备不同，配置也不同，为了达到更好的训练效率，我们需要对上述参数进行修改。

1、预训练权重--weights。我们在介绍时，采用的是yolov5s.pt预训练权重，所以得下载该权重并复制到weights文件夹下。同时，我们还得修改--weights参数，关注到这里：

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

2、模型yaml文件的路径--cfg。因为我们用的预训练模型是yolov5s.pt对应yolov5s.yaml，所以将此参数改为如下路径：

parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')

3、我们将自己的数据集文件夹改为了coco128，所以对--data参数的修改如下：

parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')

4、训练轮次。默认300，如果想更改，那么请关注到这里：

parser.add_argument('--epochs', type=int, default=300)

5、图片数量及工作核心数。这二者默认16和8。我的电脑CPU是i9-10900K(10核心，20线程)，GPU是RTX2060。如果图片输入数量设置时与核心数不匹配，那么就会出现“RuntimeError：CUDA out of memory. Tried to allocate......"的报错信息。需要修改这两个参数的同学，请关注这里：

parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')

6、避免虚拟内存不足报错

使用Pycharm编辑器进行训练时，可能会出现“OSError：[WinError 1455]页面文件太小，无法完成操作。 Error loading ......”，此时需要定位到utils文件夹下的datasets.py文件中，对第81行代码的“nw”参数进行修改，改为"num_worker=0"即可，如下图所示：

此时，即可运行train.py文件，训练自己的自定义数据集了。训练开始的截图如下图所示：

last.pt是最后一轮的权重文件，best.pt是最好的权重文件。

四、训练过程可视化

有同学注意到训练开始时，有提示说“tensorboard: Start with 'tensorboard --logdir runs/train', view at http://localhost:6006/”，所以，可以在Termianl框中输入：

tensorboard --logdir runs/train

在对应的页面看到训练过程，如下图所示：

如果模型训练完毕，那么将会使用如下命令查看训练结果：

tensorboard --logdir=runs

五、权重文件推理测验

将我们刚才训练产生的权重文件best.pt复制到weights文件夹中，使用命令行进行验证：

python detect.py --weights weights/best.pt --img 640 --conf 0.25 --source data/images/# --weights后面跟的是权重文件的路径# --conf后面是置信度# --source后面是要推理的图片路径，可以具体到文件名，还可以是整个文件夹

到对应文件查看照片即可。