DETR训练自己的数据集

1. 数据集准备

新建一个VOC2007文件夹，文件存放方式如下，其中annotations中存放.json格式的标签

2. 环境配置

（1)直接pip install -r requirements.txt

不报错，安装成功，报错见（2）

（2）分别安装requirements.txt里所需要的包

其中标红的两个，安装方式：

(1)安装pycocotools

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi
cd PythonAPI
pip install -e.

（2）安装panopticapi

git clone https://github.com/cocodataset/panopticapi.git
cd panopticapi
pip install -e.

3.预训练权重的获取

新建一个change_class_num.py的文件，代码如下：

import torch
pretrained_weights=torch.load('./detr-r50-e632da11.pth')
num_classes=4+1 # 类别数+1
pretrained_weights["model"]["class_embed.weight"].resize_(num_classes+1,256)
pretrained_weights["model"]["class_embed.bias"].resize_(num_classes+1)
torch.save(pretrained_weights,"detr_r50_%d.path"%num_classes)

运行change_class_num.py，得到一个与训练权重detr-r50_5.pth

4. 修改模型的参数

修改models/detr.py文件，build()函数中，可以将红框部分的代码都注释掉，直接设置num_classes为自己的类别数+1

5. 训练

修改main.py

根据自己的需要修改lr,epoch等

修改数据集的路径，根据自己存放的数据集路径写

设置输出路径，我自己建了个output的文件夹，写的就是文件夹的路径，修改default的值

在dataset parameters中找到修改resume为自己的预训练权重文件路径

6. 运行main.py