一 训练环境:
ubuntu 20.04 + pytorch 1.7 + torchvison 0.8
gpu: gtx1080 8GB
二 下载 yolov5 模型及权重:
https://github.com/ultralytics/yolov5
利用 weigths/download_weights.sh 来下载 yolov5s.pt , 必要的时候 科鞋上网
三 组织 训练数据
3.1 数据集下载 (感谢学长):
https://pan.baidu.com/s/1ZJbzf5qF9GKEqyo0RskWiQ
密码： 53hn
3.2 准备 yolo 训练使用的数据集格式
把第一个 红框的数据集拿来训练 ， 我整理好了 yolo系列需要的数据格式，分享一下：
https://github.com/walletiger/facemask_annotations
一图来了解 ultralytics/yolovXX 系列模型需要的数据集格式：
把网上下载的 口罩数据集 train 和 val 数据解压 到 images/train , images/val
3.3 准备 facemask.yaml 文件
放到 yolov5/data目录下， 格式如下：
1 2 3 4 5 6 7 8 | train: /workspace/data/facemark/facemask_annotations/facemask_labels_yolov5/dataset/images/train # 128 images val: /workspace/data/facemark/facemask_annotations/facemask_labels_yolov5/dataset/images/val # 128 images # number of classes nc: 2 # class names names: ['face', 'facemark'] |
四 训练 过程
1 | python3 train.py --img 640 --batch 16 --epochs 160 --data facemask.yaml --weights yolov5s.pt --cache-images |
在 runs/train/expX 下可以看到 训练 收敛过程和 验证结果
results.txt 记录了每个 epoch 训练的 精度 、召回率、mAP等数据，
train_batchx.jpg 是阶段性验证结果 ，如下示例：
五 在jetson nano 下验证模型
5.1 下载 yolov5 代码
1 | git clone https://github.com/ultralytics/yolov5 |
5.2 上传训练 好的模型：
我的在训练 主机 runs/exp3/weights/best.pt ， 放到 jetson nano yolov5/
5.3 测试一把 video 检测
1 | python3 detect.py --source /workspace/data/360p1.mp4 --weights best.pt --conf 0.3 |
准确率还可以， 速度不快 110~120 ms 一帧
据说 可以通过 转 onnx -> tensortRT 获得开挂般的性能。。 后续 再研究。