欢迎访问我的网站,希望内容对您有用,感兴趣的可以加入免费知识星球。

使用google colab训练YOLOv5模型

YOLO 迷途小书童 3年前 (2020-11-01) 16801次浏览 8个评论

软硬件环境

  • google colab
  • YOLOv5

视频看这里

此处是 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,显存是16 G。貌似每次 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

参考资料

喜欢 (5)

您必须 登录 才能发表评论!

(8)个小伙伴在吐槽
  1. 请问一下我训练时跳出CUDA out of memory是什么情况
    匿名2021-05-30 13:13
  2. 問題解決了
    匿名2021-02-24 10:00
  3. 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
  4. 你好,为什么创建时more里面没有Google colab啊?关联更多应用里也没有
    匿名2021-02-05 10:52
Ads Blocker Image Powered by Code Help Pro

Ads Blocker Detected!!!

请关闭 Adblock 等类似浏览器插件,然后刷新页面访问,感谢您的支持!

We have detected that you are using extensions to block ads. Please support us by disabling these ads blocker.