从零开始手把手教你利用yolov5训练自己的数据集(含coco128数据集/yo

利用YOLOV5训练自己的数据/

执行以下命令做好准备工作

git clone https://github.com/ultralytics/yolov5 # 下载 yolov5 项目
python3 -c "from yolov5.utils.google_utils import gdrive_download; gdrive_download('1n_oKgR81BJtqk75b00eAjdv03qVCQn2f','coco128.zip')" # 下载官方例子的数据集
cd yolov5 #进入yolov5 项目文件
pip install -U -r requirements.txt #安装需求

有的人在python3 -c 下载数据集这一步出错，现提供百度云下载地址。

另附权重文件百度云下载地址。

coco128 提取码：55q6

yolov5权重文件 密码: gt7l

权重文件更新了20200728-13:33

该项目github地址：github-yolov5

1.requirements

很多人在 pip install -U -r requirements.txt 这一步报错了，比如我，那么就请手动安装吧！

依次执行 pip install

如下方式

pip install numpy==1.17
pip install python==3.7
###下面所有的需要全部装一遍，可以参照 yolov5下的 requirements.txt 文件

Python>=3.7

PyTorch>=1.5

Cython

numpy==1.17

opencv-python

torch>=1.5

matplotlib

pillow

tensorboard

PyYAML>=5.3

torchvision

scipy

tqdm

git+#subdirectory=PythonAPI

2.训练自定义数据 2.1创建Dataset.yaml

这里官方的例子使用的是coco128.yaml，这是一个小型教程数据集，由COCO train2017 中的前128张图像组成。在此示例中，这些相同的128张图像用于训练和验证。coco128.yaml定义1）训练图像目录的路径（或带有训练图像列表的* .txt文件的路径），2）与我们的验证图像相同的路径，3）类数，4）类列表名称：

＃train和Val的数据集（图像目录或* .txt与图像路径文件）

train： …/coco128/images/train2017/

VAL： …/coco128/images/train2017/

＃class number

nc： 80

＃类名称

名称： [“人”，“自行车”你下载内容包含文件夹,需要使用百度云管家，“汽车”，“摩托车”，“飞机”，“公共汽车”，“火车”，“卡车”，“船”，“交通灯”，

“消防栓” ”，“停车标志”，“停车计时器”，“长凳”，“鸟”，“猫”，“狗”，“马”，“绵羊”，“牛”，

“大象”，“熊”，“斑马” '，‘长颈鹿’，'背包，‘雨伞’，‘手袋’，‘领带’，‘手提箱’，‘飞盘’，

‘滑雪板’，‘滑雪板’，‘运动球’，‘风筝’，‘棒球棒’，“棒球手套”，“滑板”，“冲浪板”，

“网球拍”，“瓶”，“酒杯”，“杯”，“叉”，“刀”，“勺子”，“碗”，“香蕉” ，“苹果”，

“三明治”，“橙色”，“西兰花”，“胡萝卜”，“热狗”，“比萨饼”，“甜甜圈”，“蛋糕”，“椅子”，“沙发”，

“盆栽”，“床”，“餐桌” ”你下载内容包含文件夹,需要使用百度云管家，“厕所”，“电视”，“笔记本电脑”，“鼠标”，“遥控器”，“键盘”，

“手机”，“微波炉”，“烤箱”，“烤面包机”，“水槽”，“冰箱” ，“书”，“时钟”，“花瓶”，“剪刀”，

“泰迪熊”，“吹风机”，“牙刷”]

** 如果使用自己建立的数据集 要修改的参数有

train：换成你自己的训练集地址

val：换成自己的val地址

nc: 换成你自己定义的类别的数量

类名称：换成自己定义的类别的名称

2.2创建标签

使用Labelbox和CVAT工具进行图像标记，将标签导出为darknet格式，×.txt每张图像一个文件(如果图像中没有对象，则不需要×.txt文件)。txt文件内容有：

每个图像的标签文件应该可以通过在其路径名中简单地替换/images/ *.jpg为/labels/*.txt来定位。

例如

数据集/图像/ train2017 / 000000109622.jpg ＃图像

数据集/标签/ train2017 /000000109622.txt ＃标签

具有5个人（所属类别定义为0）的标签文件

2.3组织目录

数据集/coco128放在/yolov5项目的旁边，数据集/coco128下包含两个文件夹。

一个为/coco128/images，另一个为/coco128/labels，里面分别存放数据集图片和数据集标注信息。

2.4选择模型

从./models文件夹中，选择模型。yolov5提供了五个模型，这里我们以yolov5s.yaml为例子，双击打开，更新其中参数适应2.1中定义的类别。这里一般不用动更新nc就可以了。就是匹配你自己数据集的类别就可以了，其他的参数一般不做优化不需要动。

#parameters

nc: 80

#number of classes depth_multiple: 0.33 # model depth multiple

width_multiple: 0.50 # layer channel multiple

#anchors

anchors:

- [10,13, 16,30, 33,23] # P3/8

- [30,61, 62,45, 59,119] # P4/16

- [116,90, 156,198, 373,326] # P5/32

#YOLOv5 backbone

backbone:

#[from, number, module, args]

[[-1, 1, Focus, [64, 3]], # 1-P1/2

[-1, 1, Conv, [128, 3, 2]], # 2-P2/4

[-1, 3, Bottleneck, [128]],

[-1, 1, Conv, [256, 3, 2]], # 4-P3/8

[-1, 9, BottleneckCSP, [256, False]],

[-1, 1, Conv, [512, 3, 2]], # 6-P4/16

[-1, 9, BottleneckCSP, [512, False]],

[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32

[-1, 1, SPP, [1024, [5, 9, 13]]],

[-1,12, BottleneckCSP, [1024, False]], # 10

#YOLOv5 head

head: [[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)

[-2, 1, nn.Upsample, [None, 2, ‘nearest’]],

[[-1, 6], 1, Concat,[1]], # cat backbone P4

[-1, 1, Conv, [512, 1, 1]],

[-1, 3,BottleneckCSP, [512, False]],

[-1, 1, nn.Conv2d, [na * (nc + 5), 1,1, 0]], # 16 (P4/16-medium)

[-2, 1, nn.Upsample, [None, 2, ‘nearest’]],

[[-1, 4], 1, Concat,[1]], # cat backbone P3

[-1, 1, Conv, [256, 1, 1]],

[-1, 3,BottleneckCSP, [256, False]],

[-1, 1, nn.Conv2d, [na * (nc + 5), 1,1, 0]], # 21 (P3/8-small)

[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)

2.5训练

可以从0开始训练也可以加载–cfg yolov5s.yaml --weights 通过传递匹配的权重文件从预训练的检查点进行训练：–cfg yolov5s.yaml --weights yolov5s.pt。

# Train YOLOv5s on coco128 for 5 epochs
$ python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ''

需要调整的参数

train.py

if __name__ == '__main__':
    check_git_status()
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=5)
    parser.add_argument('--batch-size', type=int, default=4)
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='*.cfg path')
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', action='store_true', help='resume training from last.pt')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--notest', action='store_true', help='only test final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    parser.add_argument('--weights', type=str, default='', help='initial weights path')
    parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--adam', action='store_true', help='use adam optimizer')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%')
    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
    opt = parser.parse_args()

epochs：指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。

batch-size：一次看完多少张图片才进行权重更新，梯度下降的mini-batch,显卡不行你就调小点。

cfg：存储模型结构的配置文件

data：存储训练、测试数据的文件

img-size：输入图片宽高,显卡不行你就调小点。

rect：进行矩形训练

resume：恢复最近保存的模型开始训练

nosave：仅保存最终checkpoint

notest：仅测试最后的epoch

evolve：进化超参数

bucket：gsutil bucket

cache-images：缓存图像以加快训练速度

weights：权重文件路径

name： 重命名results.txt to results_name.txt

device：cuda device, i.e. 0 or 0,1,2,3 or cpu

adam：使用adam优化

multi-scale：多尺度训练，img-size +/- 50%

single-cls：单类别的训练集

test.py

f __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='test.py')
    parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path')
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
    parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS')
    parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
    parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--merge', action='store_true', help='use Merge NMS')
    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
    opt = parser.parse_args()
    opt.save_json = opt.save_json or opt.data.endswith('coco.yaml')
    opt.data = check_file(opt.data)  # check file
    print(opt)

同理这里的batch_size image_size ,显卡不行你就调小点。

这就解决了yolov5 训练时会出现的pytorch出现RuntimeError: CUDA out of memory.

2.6可视化

这算是yolov5提供的一个新的功能，在训练开始之后，可以通过查看train*.jpg图像以查看训练图像，标签和增强效果。

在每代训练完成后，可以通过test_batch0_gt.jpg，来看它们的标签还真是标注的边框。

可以通过查看test_batch0_pred.jpg 来查看每代的预测结果

训练的损失和性能指标保存在tensorboard和results.txt日志文件中。训练完成后results.txt绘制results.png。在这里，我显示了从coco128开始训练到100个代YOLOv5，从零开始（橙色），从预训练的yolov5s.pt权重（蓝色）开始：

3 补充 3.1 解决 KeyError: "weights/yolov5s.pt is not compatible with models/yolov5s.yaml.

这个很明显权重文件过期了，没办法，需要重新下载，下载地址

yolov5权重下载地址。

这个是谷歌云，有点慢，开VPN吧。这两天有空的话我全下载下载把百度云更新下。

下载好了放进weights文件夹就好了。

3.2解决 运行 train 然后pytorch出现RuntimeError: CUDA out of memory.

参照2.5训练，改小参数就可以了。就可以解决了

3.3 更新

目前github还在更新，定期下载更新下yolov5文件