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为自己的预训练权重文件路径