【模型复现】自制数据集上复现刚发布的最新 yolov9 代码
Code 链接: YOLOV9
Paper 链接: YOLOV9: Learning What You Want to Learn Using Programmable Gradient Information
1. 环境安装
# 创建环境
conda create -n yolov9 python=3.7
# 激活环境
conda activate yolov9
# torch 安装
# 本机 CUDA 为 11.8,故安装了符合要求的 pytorch==1.13,这里需要自行根据 CUDA 版本安装适配的 torch 版本
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
# pip 包
cd yolov9
pip install -r requirements.txt
2. 数据准备
- 准备
yolo
格式的训练数据,数据集的文件结构为:my_datasets/ ├─train.txt ├─val.txt ├─images │ ├──train │ │ ├──xxx.jpg │ │ └──xxx.jpg │ └──val │ ├──xxx.jpg │ └──xxx.jpg └──labels ├──train │ ├──xxx.txt │ └──xxx.txt └──val ├──xxx.txt └──xxx.txt
3. 模型复现
3.1 模型文件修改
- 在
yolov9/models/detect
路径下选择yolov9-c.yaml
文件,修改nc
为训练的标签数量
3.2 数据文件修改
- 在
yolov9/data
路径下新建my_datasets.yaml
文件,以路径下的coco.yaml
为标准进行参数配置,- 修改
path
为数据存储路径 - 修改
names
为对应的标签名
- 修改
3.3 预训练权重下载
- 在
README.md
文件中点击下述链接即可下载对应的模型权重文件
- 注意:截止目前只支持 C 和 E 版本的权重下载,S 和 M 版本预训练权重尚未发布!!!
3.4 训练超参数配置
- 在
yolov9
路径下进入train_dual.py
文件进行超参数配置,常用需修改参数如下:weights
:预训练权重路径cfg
:模型配置文件路径data
:数据配置文件路径epoches
:训练epoch
数量batch-size
:训练batch-size
数量imgsz
:图片大小optimizer
:优化器
- 训练中按需进行超参数配置
3.5 模型训练
- 单卡训练指令
python train_dual.py
- 多卡训练指令
python -m torch.distributed.launch --nproc_per_node 8 train_dual.py
- 在终端中运行训练命令,若看到下述界面,即成功复现!!!
4. 模型验证
4.1 验证超参数配置
- 在
yolov9
路径下进入val_dual.py
文件进行超参数配置,常用需修改参数如下:data
:数据配置文件路径weights
:训练完成的best.pt
权重路径batch-size
:验证batch-size
数量imgsz
:图片大小task
:验证数据集选择,如val
或test
- 验证中按需进行超参数配置
4.2 训练效果验证
- 单卡验证指令
python val_dual.py
- 多卡验证指令
python -m torch.distributed.launch --nproc_per_node 8 val_dual.py
- 在终端中运行验证命令,若看到下述界面,即成功验证!!!
5. 模型推理
5.1 推理超参数配置
- 在
yolov9
路径下进入detect.py
文件进行超参数配置,常用需修改参数如下:data
:数据配置文件路径weights
:训练完成的best.pt
权重路径batch-size
:验证batch-size
数量imgsz
:图片大小task
:验证数据集选择,如val
或test
- 推理中按需进行超参数配置
5.2 模型推理
-
单卡推理指令
python detect.py
-
多卡推理指令
python -m torch.distributed.launch --nproc_per_node 8 detect.py
-
在终端中运行推理命令,若看到下述界面,即成功推理!!!
-
可进入保存路径进行推理图片查看,推理示例如下:
6. 踩坑记录
- 报错:AttributeError: ‘list’ object has no attribute ‘device’
- 解决方法:
- 修改
utils/general.py
脚本中第 903 行,将prediction = prediction[0]
修改为prediction = prediction[0][1]
即可.
- 修改