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

准备数据集

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

新建一个date.yaml文件
放置在timber文件夹中
使用自己数据集在google colab上进行yolov5模型训练_第1张图片
文件内容为:

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

nc为需要训练的类别

文件夹内的文件目录结构为下图所示
使用自己数据集在google colab上进行yolov5模型训练_第2张图片
打开google drive 新建一个文件夹用于colab进行yolov5训练,并且把timber文件夹进行压缩上传

colab操作

操作基本和jupyter notebook一致
打开colab新建一个笔记本使用自己数据集在google colab上进行yolov5模型训练_第3张图片
或者在google drive中打开使用自己数据集在google colab上进行yolov5模型训练_第4张图片
因为训练需要使用到gpu,因此将其打开
点击左上角修改,选择笔记本设置,点击gpu使用自己数据集在google colab上进行yolov5模型训练_第5张图片
点击右上角连接
使用自己数据集在google colab上进行yolov5模型训练_第6张图片
此时已经可以使用gpu

!nvidia-smi

使用自己数据集在google colab上进行yolov5模型训练_第7张图片
安装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
使用自己数据集在google colab上进行yolov5模型训练_第8张图片
接下来执行训练命令

!cd yolov5

进行训练

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

使用yolov5s的配置文件进行训练用时3.877小时
使用自己数据集在google colab上进行yolov5模型训练_第9张图片

使用自己数据集在google colab上进行yolov5模型训练_第10张图片

解决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执行
使用自己数据集在google colab上进行yolov5模型训练_第11张图片

防止文件丢失

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

你可能感兴趣的:(yolov5,深度学习,pytorch)