使用自己数据集在google colab上进行yolov5模型训练

准备数据集

标注好的数据集
需要把标签信息转换为txt格式
图片和标签文件需要分开存放到为images与labels文件夹中，并且分为三个文件夹分别用于自己测试的test文件夹，用于训练的train和valid文件夹，将其全部放在同一个文件夹中，因为训练的是木材瑕疵因此放在timber文件夹中

新建一个date.yaml文件
放置在timber文件夹中

文件内容为：

train: ../timber/train/images
val: ../timber/valid/images
nc: 1
names: ['s']

nc为需要训练的类别

文件夹内的文件目录结构为下图所示

打开google drive 新建一个文件夹用于colab进行yolov5训练，并且把timber文件夹进行压缩上传

colab操作

操作基本和jupyter notebook一致
打开colab新建一个笔记本
或者在google drive中打开
因为训练需要使用到gpu，因此将其打开
点击左上角修改，选择笔记本设置，点击gpu
点击右上角连接

此时已经可以使用gpu

!nvidia-smi

安装pytorch

pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

安装yolov5

!git clone https://github.com/ultralytics/yolov5.git
%cd /content/yolov5
!pip install -U -r requirements.txt

连接到自己的google drive

%cd ..
from google.colab import drive
drive.mount('/content/gdrive')

将自己上传的数据集复制到content文件夹下，并且解压

!cp /content/gdrive/MyDrive/Colab_Notebooks/timber.zip /content/
!unzip /content/timber.zip -d /content/timber

打开配置文件进行修改
nc为需要训练的类别改为1

接下来执行训练命令

!cd yolov5

进行训练

!python train.py --img 640 --data /content/timber/date.yaml --cfg models/yolov5s.yaml

使用yolov5s的配置文件进行训练用时3.877小时

解决Colab长时间无操作就掉线的问题

这段代码的功能是设置每隔60000ms也就是1min自动点击一下Colab的“连接”操作，这样就不至于长时间误操作而导致自动断开连接啦
在Colab界面打开控制台（快捷键Ctrl+Shift+I），复制以下代码

function ClickConnect(){
    
    
    console.log("Working");
    document.querySelector("colab-toolbar-button#connect").click()
}
setInterval(ClickConnect, 60000)

在Colab界面打开控制台（快捷键Ctrl+Shift+I），复制以上代码，到对话框中，并按enter执行

防止文件丢失

我们可以把yolov5安装在google drive中
再进行上面的操作