知识星球: 迷途小书童;个人微信:xituxiaoshutong100

使用google colab训练YOLOv5模型

ML/DL 迷途小书童 0评论

软硬件环境

视频看这里

此处是youtube的播放链接,需要科学上网。喜欢我的视频,请记得订阅我的频道,打开旁边的小铃铛,点赞并分享,感谢您的支持。

colab是什么

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

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

参考资料

喜欢 (0)
发表我的评论
取消评论

表情

Hi,您需要填写昵称和邮箱!

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址