谷歌在线训练模型的流程:
如果本地训练模型,会长时间占用自己显卡的显存,而且显卡的显存大小很大程度上影响训练的时间,谷歌提供免费的云端GPU,有足足16G的显存可以用来训练数据集。
由于是免费资源,所以训练的持续时间不能太长从而占用公共资源,如果是需要训练几天的大型数据,本教程不再适用。
1.注册Google Drive云盘
目的是为了之后将云盘中的文件挂载到Colaboratory上,防止文件丢失。
新建一个文件夹,将训练所需的准备文件放置在此文件夹中。
2.从云盘连接到Colaboratory
点击右上角新建一个Google Colaboratory,如果没有该选项,则需要点击连接更多应用程式进行安装:
创建好后,进入下面的页面:
默认是使用CPU的资源,需要更改为使用GPU的资源。点击菜单栏中的修改按钮,再点击笔记本设置,更改为使用GPU资源:
点击右上角连接按钮,选择连接到托管代码执行程序,这样就连接上云盘了:
3.准备工作:
查看当前GPU情况,16G的显存,已经算很好的配置了:
!nvidia-smi
脚本定义挂载网盘并执行代码,需要点击链接输入验证码:
import os
from google.colab import drive
drive.mount('/content/drive')
进入到挂载的文件夹中:
%cd /content/drive # /content/drive
%cd My\ Drive/colab # /content/drive/My Drive/colab
!git clone https://github.com/ultralytics/yolov5.git # 克隆yolov5的代码到文件夹中
%cd yolov5/ # /content/drive/MyDrive/colab/yolov5
回到Dirve云盘中,此时文件夹结构应该如下图,一个是yolov5官方的文件夹,一个是自己的数据集文件夹:
开始训练前,还需要修改下yolov5的文件。找到yolov5文件夹中的models文件夹,根据所需打开.yaml文件,这里以yolov5s.yaml为例:
修改其中的nc值,根据你自己数据集所需进行修改,修改完后点击右上角的Save to Drive按钮:
4.开始训练:
进入yolov5文件夹:
%cd yolov5/
执行训练命令,注意自己的训练集位置输入是否正确:
!python train.py --data ../Lion/data.yaml --cfg models/yolov5s.yaml --weights '' --batch-size 32
训练结束后,将训练好的权重中最好的权重best.pt下载到本地,其余检测操作同本地检测图片操作相同。