使用google colab训练YOLOv5模型

软硬件环境

• google colab

• YOLOv5

colab是什么

colab是google提供的一个jupyter notebook工具，支持google drive、tensorflow在内的google全家桶，主要用于机器学习的开发和研究。colab最大的好处是给广大的AI开发者提供免费的gpu资源，可以在上面非常轻松地运行如tensorflow、pytorch、keras等深度学习框架。

YOLOv5模型训练

来到google drive，点击左上方的New

colab_yolov5

创建一个新的文件夹colab

colab_yolov5

接下来，将准备好的口罩数据集上传到colab文件夹中，这个数据集，前面我们在YOLOv5模型训练的时候用过，可以到下面的地址下载

原始链接 https://public.roboflow.ai/object-detection/mask-wearing

CSDN下载 https://download.csdn.net/download/djstavaV/12624588

百度网盘 https://pan.baidu.com/s/15GSPiJ59dg4kNyUch6W5Xw， 提取码:wja4

接下来创建colab，点击New --> More --> Google Colaboratory

colab_yolov5

创建好notebook后，需要来到 修改 --> 笔记本设置 设置gpu加速

colab_yolov5

硬件加速器，选择GPU，保存

colab_yolov5

点击右上角的 连接，选择 连接到托管代码执行程序。

colab_yolov5

GPU环境设置好后，我们就可以在notebook中查看colab提供的gpu资源了，使用!nvidia-smi命令

colab_yolov5

可以看到google提供的硬件是是tesla P100，显存是16G。貌似每次colab分配的gpu是不一样的，有时候是P100，有时候是T4

下面看看pytorch的安装情况，执行

import torch
torch.__version__

可以看到平台已经默认安装，且版本是1.6，CUDA的版本是10.1

colab_yolov5

colab_yolov5

如果需要安装第三方库，可以在单元格中直接安装，如!pip3 install torchvision

colab_yolov5

接下来就把google drive挂载过来，这样就可以在colab中使用google drive中的资源了

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

path = "/content/drive/My Drive"

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

执行上述单元格中的代码，会要求进行输入验证码

colab_yolov5

colab_yolov5

google drive就会被挂载到目录/content/drive，后续就可以对google drive里的文件进行操作了

colab_yolov5

准备工作搞定好，我们就可以下载YOLOv5的源码了，在单元格中执行

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

然后切换到google drive，修改yolov5/models/yolov5s.yaml，将原来的nc: 80改为nc: 2

colab_yolov5

然后就可以来到colab，进入到yolov5目录，在单元格中执行训练命令

!python train.py --data ../mask/data.yaml --cfg models/yolov5s.yaml --weights '' --batch-size 64

训练时，出现了pyyaml模块的一个错误，这是由于pyyaml版本过低的原因，我们升级下就可以解决

pip install -U pyyaml

colab_yolov5

继续训练

colab_yolov5

colab_yolov5

参考资料

• https://github.com/ultralytics/yolov5

• https://xugaoxiang.com/2020/07/02/yolov5-training/

• https://xugaoxiang.com/2020/06/17/yolov5/