使用google colab训练YOLOv5模型
source link: https://xugaoxiang.com/2020/11/01/google-colab-yolov5/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
软硬件环境
- google colab
- YOLOv5
视频看这里
此处是youtube
的播放链接,需要科学上网。喜欢我的视频,请记得订阅我的频道,打开旁边的小铃铛,点赞并分享,感谢您的支持。
colab是什么
colab
是google
提供的一个jupyter notebook
工具,支持google drive
、tensorflow
在内的google
全家桶,主要用于机器学习的开发和研究。colab
最大的好处是给广大的AI
开发者提供免费的gpu
资源,可以在上面非常轻松地运行如tensorflow
、pytorch
、keras
等深度学习框架。
YOLOv5模型训练
来到google drive
,点击左上方的New
创建一个新的文件夹colab
接下来,将准备好的口罩数据集上传到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
创建好notebook
后,需要来到 修改 --> 笔记本设置 设置gpu
加速
硬件加速器,选择GPU
,保存
点击右上角的 连接,选择 连接到托管代码执行程序。
GPU
环境设置好后,我们就可以在notebook
中查看colab
提供的gpu
资源了,使用!nvidia-smi
命令
可以看到google
提供的硬件是是tesla P100
,显存是16G
。貌似每次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
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK