YOLOv7训练自己的数据集
1. 准备深度学习环境
本人的笔记本环境 Ubuntu16.04, cuda-8.0, conda-4.3.8
1.1 代码克隆
git clone https://github.com/WongKinYiu/yolov7.git
1.2 环境安装
## 创建conda环境
conda create -n yolov7 python=3.6
## 安装包
source activate yolov7
## 使用 conda 单独安装 torch 和 torchvision
conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=10.2 -c pytorch
## 打开 requirements.txt, 注释掉 torch 和 torchvision, 使用 pip 安装其他工具包
# torch>=1.7.0,!=1.12.0
# torchvision>=0.8.1,!=0.13.0
pip install -r requirements.txt
1.3 权重下载
mkdir weights
cd weights
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
1.4 环境测试
## CPU测试
python detect.py --weights models/yolov7-tiny.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg --device cpu
## GPU测试
python detect.py --weights models/yolov7-tiny.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg --device 0
2. 准备数据
2.1 yolo格式的数据集准备
- 使用labelImg工具制作YOLO目标检测数据集
-
一张图片对应一个标签文件:a.jpg -> a.txt。其中a.txt文件中内容如下:
## class x y w h 0 0.421094 0.512500 0.051562 0.091667 0 0.514062 0.508333 0.034375 0.087500 0 0.591406 0.551042 0.082812 0.039583
-
数据保存目录
- /data/YOLOLable/ - train/ - images/ - a.jpg - ...... - labels/ - a.txt - ...... - val/ - images/ - aa.jpg - ...... - labels/ - aa.txt - ...... - test/ - images/ - aaa.jpg - ...... - labels/ - aaa.txt - ......
-
3. 配置训练的相关文件
3.1 修改数据配置文件
- 在 yolov7/data 目录下创建test.yaml (yolov7/data/test.yaml)
# train and val data as train: ./datasets/train.txt val: ./datasets/val.txt test: ./datasets/test.txt # number of classes nc: 1 # class names names: [ 'book']
- ./datasets/train.txt:每一行为训练图片的绝对路径
/data/YOLOLable/train/images/a.jpg .......
3.2 修改模型配置文件
- 在 yolov7/cfg/training 目录下复制yolov7.yaml,并命名为yolov7-test.yaml (yolov7/data/yolov7-test),修改 nc 的值
# parameters # nc: 80 # number of classes nc: 1 # 你的数据集中类别的个数 ...
4. 模型训练和测试
模型训练和测试可以根据yolov7根目录中 readme.md 文件描述的进行。
4.1 模型训练
## p5 models
python train.py --weights weights/yolov7.pt --data data/test.yaml --epochs 100 --batch-size 8 --cfg cfg/training/yolov7-test.yaml --workers 0 --device 0 --img-size 640 640
4.2 模型测试
python test.py --data data/test.yaml --img 640 --batch 8 --conf 0.001 --iou 0.65 --device 0 --weights runs/train/exp/weights/best.pt