yolov9训练自己的数据—vehicle 4类

yolov9训练自己的数据

[1 conda环境](#1 conda环境)
- 安装指定版本torch
[2 预训练模型测试](#2 预训练模型测试)
[3 训练自己的数据集](#3 训练自己的数据集)
- [3.1 制作数据](#3.1 制作数据)
- [3.2 创建模型配置文件](#3.2 创建模型配置文件)
- [3.3 创建数据加载配置文件](#3.3 创建数据加载配置文件)
- [3.4 使用ClearML跟踪训练日志](#3.4 使用ClearML跟踪训练日志)
- [3.5 训练](#3.5 训练)
- [3.6 模型测试](#3.6 模型测试)
- [3.7 转换成TensorRT模型](#3.7 转换成TensorRT模型)
[4 参考文档](#4 参考文档)

1 conda环境

下载yolov9代码，并执行以下命令

bash 复制代码

$ git clone https://github.com/WongKinYiu/yolov9.git
$ cd yolov9
$ conda create --name yolov9 python=3.8
$ pip install -r requirement.txt

安装指定版本torch

到pytorch官网下载安装。

bash 复制代码

pip install torch==2.0.0 torchvision==0.15.1 -index-url https://download.pytorch.org/whl/cu118

torch安装失败，到上面地址手动下载v2.0.0，再安装。

bash 复制代码

 $ pip install torch-2.0.0+cu118-cp38-cp38-linux_x86_64.whl
 $ pip install torchvision-0.15.1+cu118-cp38-cp38-linux_x86_64.whl

2 预训练模型测试

bash 复制代码

# inference converted yolov9 models
$ python detect.py --source ./data/images/horses.jpg --img 640 --device 0 --weights ./weights/yolov9-c-converted.pt --conf 0.1

# inference yolov9 models
# python detect_dual.py --source ./data/images/horses.jpg --img 640 --device 0 --weights ./weights/yolov9-c.pt

# inference gelan models
# python detect.py --source ./data/images/horses.jpg --img 640 --device 0 --weights ./weights/gelan-c.pt

bash 复制代码

detect: weights=['./weights/yolov9-c-converted.pt'], source=./data/images/horses.jpg, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.4e, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=run_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5 🚀 v0.1-35-g9660d12 Python-3.8.18 torch-2.0.0+cu118 CUDA:0 (NVIDIA GeForce RTX 3090, 24260MiB)

Fusing layers...
gelan-c summary: 387 layers, 25288768 parameters, 64944 gradients, 102.1 GFLOPs
image 1/1 /home/hard_disk/waf/projects-test/yolov9/data/images/horses.jpg: 448x640 5 horses, 89.3ms
Speed: 0.2ms pre-process, 89.3ms inference, 40.3ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs/detect/exp

3 训练自己的数据集

3.1 制作数据

标注好的数据划分成train和val，调整成以下目录结构。images内放图片，labels放对应标签，图片名和标签名要一致。

bash 复制代码

vehicle/
├── train
│   ├── images
│   │   ├── 000000000064.jpg
│   ├── labels
│   │   ├── 000000000064.txt
├── val
│   ├── images
│   │   ├── 00000000724.jpg
│   ├── labels
│   │   ├── 000000000724.txt
├── train.txt
├── val.txt

train.txt内容为图像完整路径。

bash 复制代码

/home/hard_disk/waf/data/vehicle/train/images/000000000064.jpg
/home/hard_disk/waf/data/vehicle/train/images/000000000071.jpg
...

3.2 创建模型配置文件

在yolov9/models/detect/目录下创建yolov9-c_vehicle.yaml，拷贝yolov9-c.yaml内容，并修改nc值为待训练的类别数。

3.3 创建数据加载配置文件

在yolov9/data/目录下创建vehicle.yaml，内容为

bash 复制代码

path: /home/hard_disk/waf/data/vehicle/  # dataset root dir
train: train.txt  # train images (relative to 'path') 20960 images,19335+1625(uav)
val: val.txt  # val images (relative to 'path') 100 images
test: 

# Classes
names:
  0: car
  1: bus
  2: truck
  3: train

3.4 使用ClearML跟踪训练日志

（1）ClearML官网注册账号。

（2）右上角进入个人主页，依次点"Setting"--"Workspace"--"+Creat new credentials"，复制这段api：

（3）终端安装clearml

bash 复制代码

pip install clearml -i https://pypi.tuna.tsinghua.edu.cn/simple/

运行 clearml-init 连接ClearML 服务器，按照提示粘贴之前复制的api。

3.5 训练

bash 复制代码

#gpu单卡训练
$ python train_dual.py --workers 8 --device 0 --batch 16 --data data/vehicle.yaml --img 640 --cfg models/detect/yolov9-c_vehicle.yaml --weights '' --name yolov9-c --hyp data/hyps/hyp.scratch-high.yaml --min-items 0 --epochs 500 --close-mosaic 15 --cache

3.6 模型测试

bash 复制代码

python detect.py --source E:\DJI_20240307133038_0001.MP4 --img 640 --device 0 --weights weights\yolov9-c_vehicle.pt

3.7 转换成TensorRT模型

（1）Re-parameterization，将 reparameterize.py放到yolov9目录下，perform re-parameterization:

bash 复制代码

python reparameterize.py weights\yolov9-c_vehicle.pt weights\yolov9-c_vehicle_converted_cuda.pt

reparameterize.py代码如下，nc修改为训练的类别数，device也可以设置为cpu

bash 复制代码

import argparse
import torch
from models.yolo import Model

def main(input_model_path, output_model_path):
    device = torch.device("cuda")
    cfg = "./models/detect/gelan-c.yaml"
    model = Model(cfg, ch=3, nc=4, anchors=3)
    #model = model.half()
    model = model.to(device)
    _ = model.eval()
    ckpt = torch.load(input_model_path, map_location='cuda')
    model.names = ckpt['model'].names
    model.nc = ckpt['model'].nc


    idx = 0
    for k, v in model.state_dict().items():
        if "model.{}.".format(idx) in k:
            if idx < 22:
                kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.cv2.".format(idx) in k:
                kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.cv3.".format(idx) in k:
                kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.dfl.".format(idx) in k:
                kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
        else:
            while True:
                idx += 1
                if "model.{}.".format(idx) in k:
                    break
            if idx < 22:
                kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.cv2.".format(idx) in k:
                kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.cv3.".format(idx) in k:
                kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
            elif "model.{}.dfl.".format(idx) in k:
                kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))
                model.state_dict()[k] -= model.state_dict()[k]
                model.state_dict()[k] += ckpt['model'].state_dict()[kr]
    _ = model.eval()


    m_ckpt = {'model': model.half(),
              'optimizer': None,
              'best_fitness': None,
              'ema': None,
              'updates': None,
              'opt': None,
              'git': None,
              'date': None,
              'epoch': -1}
    torch.save(m_ckpt, output_model_path)
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to convert YOLO model")
    parser.add_argument("input_model_path", type=str, help="Path to the input YOLO model")
    parser.add_argument("output_model_path", type=str, help="Path to save the converted YOLO model")
    args = parser.parse_args()
    main(args.input_model_path, args.output_model_path)

（2）导出onnx模型，在yolov9目录下执行以下命令，生成yolov9-c_vehicle_converted_cuda.onnx

bash 复制代码

python export.py --weights  weights\yolov9-c_vehicle_converted_cuda.pt --include onnx

（3）Build a TensorRT engine:

bash 复制代码

trtexec.exe --onnx=yolov9-c_vehicle_converted_cuda.onnx --explicitBatch --saveEngine=yolov9-c_vehicle_converted_cuda.engine --fp16

4 参考文档

(1) yolov9 训练自己数据集＜日志＞

(2) How to Train YOLOv9 on a Custom Dataset

(3) Convert YOLOv9-C

(4) spacewalk01/TensorRT-YOLOv9

(5) yolov5/tutorials/train_custom_data

(6) RuntimeError: The size of tensor a (3) must match the size of tensor b(64) at non-singleton dimension 1