RT-DETR(Pytorch版)训练自己coco格式数据集

  1. 从github上下载代码,链接https://github.com/lyuwenyu/RT-DETR,这里主要介绍如何使用pytorch版训练

  2. 将自己coco格式据集放到目录rtdert_pytorch下;此外数据集目标下还需要classes.txt 文件,内容如下:

person
bicycle
plane
  1. 根据自己数据集,修改rtdetr_pytorch/configs/dataset/coco_detection.yml文件中的
num_classes: 80

以及train_dataloader与val_dataloader

    img_folder: ./dataset/coco/train2017/
    ann_file: ./dataset/coco/annotations/instances_train2017.json

5 修改rtdetr_pytorch/src/data/coco/coco_dataset.py的names和label2category,我的数据集有两个类别,json文件的类别id分别为0和1;对应的修改为如下:

names = {
    1: 'person',
    2: 'plane'
}

label2category = {
    1: 0,
    2: 1,
}

6 修改文件rt_detr_pytorch/src/zoo/rtdetr/rtdetr_postprocessor.py中的函数mscoco_label_category_map(self, )为如下:

    def mscoco_label_category_map(self, ):
        return {
            1: 0,
            2: 1
        }
       

具体的类别需要根据自己的数据集调整。

你可能感兴趣的:(pytorch,人工智能,python)