colab训练yolov5自己的数据集

这篇博文是复现OW-DETR时遇到了困难，然后想到了我跑yolo时的种种困难，把跑通的过程给大家写出来供大家参考，我在云上跑的主要原因是自己的笔记本没有gpu，但是又不想去办公室，感觉在云上比较方便，但其实困难很多。

yolov5源码：ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)

1.第一步仍然是把google drive挂载过来

import os
from google.colab import drive
drive.mount('/content/drive')

path = "/content/drive/My Drive"

os.chdir(path)
os.listdir(path)

2.下载yolov5的源码

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

这样就可以在MyDive中看到下载好的yolov5文件夹

3.安装yolov5需要的依赖库

!pip install -r ./yolov5/requirements.txt

大部分库都是安装好的，但是还有有几个库是新安装的。

4.更改为gpu运行

在修改，笔记本设置中，改成gpu运行。

5.跑一遍源码中的coco128

!python /content/drive/MyDrive/yolov5/train.py --data /content/drive/MyDrive/yolov5/data/coco128.yaml --cfg /content/drive/MyDrive/yolov5/models/yolov5s.yaml --weights 'yolov5s.pt' --batch-size 64

程序会根据yaml文件自己下载解压coco128数据集以及yolov5s.pt。

然后就开始训练一百轮，大概需要十来分钟。可以看到训练结果保存在了这个文件夹中。

在文件夹中可以看到训练的各个参数的变化以及train和val的各个可视化数据，在weight文件夹中会生成两个权重文件，best是效果最好的，last是最后一次训练的。

6.预测

其实detect.py是不需要我们自己训练coco128的权重可以直接运行的，但是这里我想用自己跑出来的这个权重文件来进行预测看效果。

!python /content/drive/MyDrive/yolov5/detect.py --weights /content/drive/MyDrive/yolov5/runs/train/exp/weights/best.pt --data ./data/coco.yaml --img 640

我这里已经有了exp文件夹，所以他保存在了exp2文件夹中，可以看到预测完的两个图片，我这里就放了一个图片。

7.制作自己的数据集

制作自己的数据集这又是一个很大的工作，大家可以参考网上的教程，用labelimg进行标注，但是我的labelimg在标注时会闪退，我又偷懒没想去解决这个问题，再加上我们组正在做自己的数据集，我后面要做开放世界检测，现在标注这个事我暂时用不到，所以我就偷了个懒，在网上找了现成的数据集，大家如果不想自己标注，只想找一些现成的数据集进行练习的话，可以去roboflow上下载一些现成的yolov5格式的数据集，当然roboflow也可以进行一些标注工作，但是我用起来可能由于网络原因，比较慢。

8.把数据集上传到谷歌云盘中

这一步也是非常浪费时间的一步，这一步真的非常慢，这也是colab不方便的原因之一，需要注意的一点是，要把数据集上传到yolov5文件夹中，singlanguage.v5i.yolov5pytorch这个是我上传的自己的数据集。

9.用自己的数据集进行训练

训练自己的数据集和用源码的coco128数据集差不多，就是多了一步配置yaml文件

配置自己的yaml文件

打开yolov5文件夹下的data文件夹，创建一个自己的yaml文件，我这里是sing.yaml文件，格式可以按照coco128.yaml文件进行，nc改成自己的类别数，name改成自己的类别名。

如果是在roboflow中下载的数据集，数据集文件夹下有一个data.yaml文件，可以直接把里面的内容复制到data文件夹下的yaml文件中，然后我是将train和val的路径补全了，我忘记了不补全可不可以用，还有就是要把文件夹名称中的空格都删掉，路径中不能有中文和空格。

开始训练

!python /content/drive/MyDrive/yolov5/train.py --data /content/drive/MyDrive/yolov5/data/sing.yaml --cfg /content/drive/MyDrive/yolov5/models/yolov5s.yaml --weights '' --batch-size 64

代码的话就把--data改成自己的yaml文件路径，图片多的话运行就会比较慢，我当时选了一个比较大的数据集跑了一百轮，跑了挺长时间，如果数据集比较大，可以先把轮数调小，在代码后加上    --epochs +自己的轮数。

我这里用了两个小时，训练完成后结果保存在了exp2这个文件夹中，在weights文件夹中也可以看到两个权重文件，之后就可以进行预测，预测的方法和上面利用coco128训练后预测的方法一致。

我这里偷懒没有预测，只是看了train文件夹下的val_batch0_pred文件，这是在验证集上验证的效果。

在云端训练yolov5自己的数据集基本就是这些内容，大家可以尝试看一下源码，yolov5的源码写的还是很清晰的。