折腾了大概一个星期, 最大的问题就是python的版本的问题.
最重要感谢小蒋提示了我, 转模型要3.6, 训练要3.8, 顺这个这思路走.
首先, 需要一个ubuntu 20的虚拟机或者物理机.
然后第一时间装上conda, 装mini的就行.
https://repo.anaconda.com/miniconda/Miniconda3-py38_4.12.0-Linux-x86_64.sh
在ubuntu终端运行
$ bash Miniconda3-py38_4.12.0-Linux-x86_64.sh
一步步往下走就行, 中间什么用户协议输入yes
是否初始化选择yes
关掉终端, 重新打开.
跟windows版本一样, 终端前面可以看到conda目前的环境.
$ conda create -n train_yolo_py38 python=3.8
激活py38环境
$ conda activate train_yolo_py38
$ git clone https://github.com/airockchip/yolov5.git
修改data/下面的yaml文件, 就像你平时做的那样, 训练模型.
如果是用显卡进行训练的, 可以考虑安装windows版本的conda跟python3.8 以及对应的依赖, 然后在windows下面训练模型, 因为显卡没法挂载到虚拟机里面.
比如我想训练一个安全帽的模型
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/safe_hat # dataset root dir
train: images/train # train images (relative to 'path') 128 images
val: images/val # val images (relative to 'path') 128 images
test: images/test # test images (optional)
# Classes
nc: 2 # number of classes
names: ['person', 'hat'] # class names
用yolov5的训练方法训练, 这里生成的权重/模型, 其实已经被优化过了, 一会儿看pt文件就知道, silu已经被换成了relu
(py38torchgpu) D:\rkai\yolov5>python train.py --weights '' --cfg yolov5s.yaml --data safe_hat.yaml --epochs 300 --batch-size 16
修改一下batch-size, 使得显卡内存利用率更高, 让训练速度更快, 我是8G的现存, batch-size设置成32.
训练的目标就是让[email protected] 跟[email protected]:.95尽量的高
每个Epoch即整个训练集, 约4000张图片, 每次32张图片(batch-size), 进行一次前向传播, 再用损失函数, 反向传播, 求偏导数, 这样就是一个epoch, 暂时定为重复300个epoch, 如果中间程序发现长时间mAP提升, 会提前停止训练.
50轮的时候, mAP已经是0.914了, 因为我的val跟train的样本集是分开的, 所以不会出现过拟合的问题.
生成的权重文件, 就在runs\train\expX\weights\best.pt
X是个会自增的数字, 每训练一次都会加1
训练成果图片如上
直接拿python脚本测试一下
$ python detect.py --weights runs\train\exp4\weights\best.pt --source D:\\rkai\\datasets\\safe_hat\\images\\test\\person_hat_2_72.jpg