这是我参与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
接下来,将准备好的口罩数据集上传到 colab 文件夹中,这个数据集,前面我们在 YOLOv5 模型训练的时候用过,可以到下面的地址下载
原始链接 public.roboflow.ai/object-dete…
CSDN 下载 download.csdn.net/download/dj…
百度网盘 pan.baidu.com/s/15GSPiJ59…, 提取码: wja4
接下来创建 colab,点击 New --> More --> Google Colaboratory
创建好 notebook 后,需要来到 修改 --> 笔记本设置 设置 gpu 加速
硬件加速器,选择 GPU,保存
点击右上角的 连接,选择 连接到托管代码执行程序
GPU 环境设置好后,我们就可以在 notebook 中查看 colab 提供的 gpu 资源了,使用 !nvidia-smi 命令
可以看到 google 提供的硬件是 tesla P100,显存是16 G。不过貌似每次 colab 分配的 gpu 是不一样的,有时候是 P100,有时候是 T4,对于大部分的应用都是够用的
下面看看 pytorch 的安装情况,执行
import torch
torch.__version__
可以看到平台已经默认安装,且版本是1.6,CUDA 的版本是10.1
如果需要安装第三方库,可以在单元格中直接安装,如 !pip3 install torchvision
接下来就把 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)
执行上述单元格中的代码,会要求进行输入验证码
google drive 就会被挂载到目录 /content/drive,后续就可以对 google drive 里的文件进行操作了
准备工作搞定好,我们就可以下载 YOLOv5 的源码了,在单元格中执行
!git clone https://github.com/ultralytics/yolov5.git
然后切换到 google drive,修改 yolov5/models/yolov5s.yaml,将原来的 nc: 80 改为 nc: 2
然后就可以来到 colab,进入到 yolov5 目录,在单元格中执行训练命令
!python train.py --data ../mask/data.yaml --cfg models/yolov5s.yaml --weights '' --batch-size 64
训练时,出现了 pyyaml 模块的一个错误,这是由于 pyyaml 版本过低的原因,我们升级下就可以解决
pip install -U pyyaml
继续训练
搞定!