【colab】在colab上使用yolov5训练自己的模型

一.开始

本文记录了在谷歌平台colab上使用yolov5训练自己的模型，包括具体流程和ipynb代码。使用前按照自己的需要，选择开通Google Drive云盘空间（我是开通了200G），以及Colab Pro+。使用colab的优点是价格较其他云平台较实惠，缺点是虚拟机不固定，每隔一段时间有人机验证，每次最多使用24小时（普通版是12小时）

二.修改并上传项目文件

出于方便修改的原因，我是在本地先下载yolov5项目，修改好具体代码，然后再上传到Google drive的。当然也可以在colab/drive中修改代码

1. 下载项目

到yolov5官网下载项目文件，本次项目选用了最新的V5.0版本

2. 制作yolo数据集

我首先是使用了ImgLabel标注自己的图像，制作原始数据集（由于之前做的项目是mmdet，所以不是yolo格式，是xml格式的）；之后使用代码将xml数据集转换成YOLO格式即可；如图修改classes为自己项目的种类，TRAIN_RATIO是训练集的比例，这里设置成80%（测试集就为20%）

• 创建一个名为VOCdevkit文件夹，其子文件夹如下
  VOCdevkit
  ——VOC2007
  ——ab.yaml（训练使用的文件）
  ————Annotations（放xml文件）
  ————JPEGImages（放jpg图片）
  xml_yolo.py（和VOCdevkit同级目录）
  之后运行xml_yolo即可得到yolo格式的数据集

ab.yaml的内容如下

train: ../VOCdevkit/images/train
val: ../VOCdevkit/images/val
nc: 3
names: [ 'classes1','classes2','classes3']

如果使用v6.0的话可能弹出找不到数据集的话，则使用下方代码

train: /content/data/VOCdevkit/images/train
val: /content/data/VOCdevkit/images/val
nc: 3
names: [ 'classes1','classes2','classes3']

为了方便在Colab上训练，VOCdevkit压缩成zip文件，上传到Google Drive当中（因为colab训练yolov5时候，需要先将数据集从Drive上传到虚拟机当中，如果数据集太大，第一轮训练会耗费特别多时间）；

3. 修改参数

修改参数方便就是按照自己的要求来进行修改，主要修改class_nums（model/下的yaml文件）为自己的种类，然后修改train.py（红色也必须修改的，绿色的按自己需求修改）；其他参数可以参照网络上其他博客

4. 上传项目

在Google Drive中先创建一个空的文件夹命名为yolov5（或者其他），之后在该文件下上传(Folder upload)整个项目（yolov5-master），然后推荐再创建一个文件夹（我自己命名为save）用于保存训练好的权值，这样下次训练速度也比较快（第一轮就少加载之前的权值了）

至此项目修改完成，准备到colab中写代码运行项目

三.新建colab笔记本（yolov5.ipynb）

打开Colab，首先新建yolov5.ipynb笔记本，然后打开 ①修改，将笔记本设置为GPU模型，有条件升级为高RAM（Colba Pro功能），之后点击 ②连接 到colab的虚拟机，然后开始输入代码进行训练，一般colab pro有更大的可能分到P100的显卡，普通用户可能是T4显卡或者K80。

笔记本中输入的代码如下（点击左上角+代码开始输入）：

1. 查看GPU

!nvidia-smi

2. 虚拟机上创建data文件夹，用于解压数据集

!mkdir ./data

3. 连接到自己的Google Drive

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

4. 复制并解压数据集到data里面

!cp /content/drive/MyDrive/yolov5/VOCdevkit.zip ./data/
!cd ./data && unzip VOCdevkit.zip

5. 进入到yolov5根目录下

%cd /content/drive/
%cd My\ Drive/yolov5
%cd yolov5-master/

6. 训练

!python train.py

7. 中止后继续训练

如果训练中断，例如虚拟机到24小时自动断开，或者没有及时进行人机认证，都会导致断开虚拟机，这个时候需要重新进行1-5步，然后跳过第6步，执行第7步继续训练（-- auto-resume）

!python train.py --resume

四.结语

至此在colab使用yolov5训练模型的过程就结束了