加微信进交流群:xituxiaoshutong100

使用google colab训练YOLOv5模型

ML/DL 迷途小书童 6评论

软硬件环境

视频看这里

此处是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

参考资料

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

表情
(6)个小伙伴在吐槽
  1. 問題解決了
    匿名2021-02-24 10:00 回复
    • 好的
      迷途小书童2021-02-25 13:32 回复
  2. github: up to date with https://github.com/ultralytics/yolov5 ✅ Traceback (most recent call last): File "train.py", line 479, in check_requirements() File "/content/drive/My Drive/colab/yolov5/utils/general.py", line 92, in check_requirements pkg_resources.require(requirements) # DistributionNotFound or VersionConflict exception if requirements not met File "/usr/local/lib/python3.6/dist-packages/pkg_resources/__init__.py", line 886, in require needed = self.resolve(parse_requirements(requirements)) File "/usr/local/lib/python3.6/dist-packages/pkg_resources/__init__.py", line 772, in resolve raise DistributionNotFound(req, requirers) pkg_resources.DistributionNotFound: The 'thop' distribution was not found and is required by the application 照上面做 出現這樣的錯誤
    匿名2021-02-23 15:53 回复
    • 依赖没装好。
      迷途小书童2021-02-24 17:11 回复
  3. 你好,为什么创建时more里面没有Google colab啊?关联更多应用里也没有
    匿名2021-02-05 10:52 回复
    • 肯定有,仔细找找
      迷途小书童2021-02-08 14:21 回复