DEIM模型训练自己的数据集记录

原创已于 2025-07-31 09:34:28 修改 · 1.1k 阅读

24 ·

CC 4.0 BY-SA版权

文章标签：

#计算机视觉 #python #目标检测

于 2025-07-30 16:46:06 首次发布

DEIM模型训练自己的数据集记录

1、项目拉取，预训练模型下载

github: https://github.com/Intellindust-AI-Lab/DEIM

预训练模型选择：根据需要选择基于D-FINE或RT-DERT的DEIM模型，下面一定要修改、使用对应的模型配置文件，比如下载ｄｅｉｍ－ｄ－ｆｉｎｅ－ｓ模型，步骤３就要修改数据配置文件ｄｆｉｎｅ－ｈｇｎｅｔｖ２－ｓ－ｃｏｃｏ．ｙｍｌ
在这里插入图片描述

2、配置环境

我的习惯：先创建一个虚拟环境，安装对应cuda 版本的pytorch-gpu（pytorch版本参考requirement.txt文件中要求版本），再安装其他依赖

conda create -n deim python=3.11.9
conda activate deim
pip install -r requirements.txt  (如果单独安装了pytorch，需要删除文件中的torch相关)

3、自定义数据

1）将自己的数据转换成要求的标注格式
在这里插入图片描述我的数据标注格式是yolo格式（每个图片对应一个同名的txt文件），下面是yolo转换成coco格式的脚本：

# -*- coding:utf-8 -*-
import os
import json
from tqdm import tqdm
import cv2

# 类别名
CLASS_LIST = ['B1', 'B2', 'B3', 'A1', 'A2', 'A3', 'A4', 'A5', 'A6', 'A7', 'A8', 'A9', 'A10', 'A11']

def yolo_to_coco(image_dir, txt_dir, output_json_path):
    images_dir = os.path.join(image_dir)
    labels_dir = os.path.join(txt_dir)

    if not os.path.exists(images_dir) or not os.path.exists(labels_dir):
        raise FileNotFoundError("images or labels not exist")

    categories = [{"id": i, "name": name} for i, name in enumerate(CLASS_LIST)]

    coco_data = {
        "images": [],
        "annotations": [],
        "categories": categories
    }

    image_id = 1
    annotation_id = 1
    loss_label_num = 0

    class_dict = {}
    for i in range(len(CLASS_LIST)):
        class_dict[i] = 0

    for image_name in tqdm(os.listdir(images_dir)):
        if not image_name.endswith(('.jpg', '.png', '.jpeg')):
            print(f"{image_name} not endwith jpg or png")
            continue

        image_path = os.path.join(images_dir, image_name)
        image = cv2.imread(image_path)
        height, width, _ = image.shape

        coco_data["images"].append({
            "id": image_id,
            "file_name": image_name,
            "width": width,
            "height": height
        })

        label_name = os.path.splitext(image_name)[0] + '.txt'
        label_path = os.path.join(labels_dir, label_name)
        if not os.path.exists(label_path):
            loss_label_num += 1
            image_id += 1
            print(f"{image_name} not label")
            continue

        with open(label_path, 'r') as f:
            lines = f.readlines()

        for line in lines:
            parts = line.strip().split()
            if len(parts) != 5:
                continue

            class_id, x_center, y_center, bbox_width, bbox_height = map(float, parts)
            class_id = int(class_id)
            # new_class_id = CLASS_MAPPING.get(int(class_id), None)
            if class_id is None:
                continue

            x_center *= width
            y_center *= height
            bbox_width *= width
            bbox_height *= height

            x_min = ｍａｘ(x_center - (bbox_width / 2),０)
            y_min = ｍａｘ(y_center - (bbox_height / 2),０)

            area = bbox_width * bbox_height

            coco_data["annotations"].append({
                "id": annotation_id,
                "image_id": image_id,
                "category_id": class_id,
                "bbox": [x_min, y_min, bbox_width, bbox_height],
                "area": area,
                "iscrowd": 0
            })

            class_dict[class_id] += 1

            annotation_id += 1

        image_id += 1

    with open(output_json_path, 'w') as f:
        json.dump(coco_data, f, indent=4)

    print(f"end saved to {output_json_path},loss label num:{loss_label_num}")
    print(class_dict)


def main(image_dir, txt_dir, output_json_dir):
    os.makedirs(output_json_dir, exist_ok=True)
    sub_dirs = os.listdir(txt_dir)
    for subfolder in sub_dirs:
        yolo_image_dir = os.path.join(image_dir, subfolder)
        yolo_txt_dir = os.path.join(txt_dir, subfolder)
        output_json = os.path.join(output_json_dir, f"{subfolder}.json")
        yolo_to_coco(yolo_image_dir, yolo_txt_dir, output_json)


if __name__ == "__main__":
    image_dir = r".\00labeled\zhenzhi20241130\images"
    txt_dir = r".\00labeled\zhenzhi20241130\labels"
    output_dir = r".\00labeled\zhenzhi20241130\annotations_14class"
    main(image_dir, txt_dir, output_dir)

2）修改数据配置文件（数据路径、标签类别，模型使用的数据配置文件路径）
在这里插入图片描述

4、修改配置参数：batch_size, epochs，resize_size

在这里插入图片描述

5、训练

在这里插入图片描述

6、推断

在这里插入图片描述

7、部署（onnx、tensorrt)

＃　导出ｏｎｎｘ模型
python tools/deployment/export_onnx.py --check -c configs/deim_dfine/deim_hgnetv2_${model}_coco.yml -r model.pth
＃　导出ｔｅｎｓｏｒｒｔ模型
trtexec --onnx=best_stg2.onnx --saveEngine=best_stg2.engine --buildOnly --minShapes=images:1x3x640x640 --optShapes=images:16x3x640x640 --maxShapes=images:32x3x640x640  --workspace=1024 --fp16

问题：
目前按照上述方法导出的ｏｎｎｘ模型与．ｐｔ模型结果一致，ｔｅｎｓｏｒｒｔ模型与ｏｎｎｘ结果不一致，ｔｅｎｓｏｒｒｔ检测框超级多，ＧＰＴ分析原因可能是ｎｍｓ模块没有导出到ｔｅｎｓｏｒｒｔ，遇到同样问题并解决的小伙伴请留言告知，感谢～

问题记录：

1、安装tensorrt之后，训练时出现下面的报错：

Could not load library libcudnn_cnn_train.so.8. Error: /usr/local/cuda-12.1/lib64/libcudnn_cnn_train.so.8: undefined symbol: _ZN5cudnn3cnn5infer22queryClusterPropertiesERPhS3_, version libcudnn_cnn_infer.so.8