使用Google colab训练YOLOv5

这是我参与11月更文挑战的第18天，活动详情查看：2021最后一次更文挑战

环境

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 模型训练的时候用过，可以到下面的地址下载

原始链接 public.roboflow.ai/object-dete…

CSDN 下载 download.csdn.net/download/dj…

百度网盘 pan.baidu.com/s/15GSPiJ59…，提取码: 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，显存是16 G。不过貌似每次 colab 分配的 gpu 是不一样的，有时候是 P100，有时候是 T4，对于大部分的应用都是够用的

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

import torch
torch.__version__
复制代码

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

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

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

搞定！