End-to-End Object Detection with Transformers代码实现

End-to-End Object Detection with Transformers
github：https://github.com/facebookresearch/detr

环境依赖

conda install cython scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

改类别
models/detr.py

改权值

import torch
pretrained_weights  = torch.load('detr-r50-e632da11.pth')
num_class = 7  # 这里是你自己数据集的类别数量 + 1
pretrained_weights["model"]["class_embed.weight"].resize_(num_class+1, 256)
pretrained_weights["model"]["class_embed.bias"].resize_(num_class+1)
torch.save(pretrained_weights, "detr-r50_%d.pth"%num_class)

命令 python main.py --dataset_file “coco” --coco_path data/coco --epochs
100 --lr=1e-4 --batch_size=2 --num_workers=4 --output_dir=“outputs”
–resume=“detr-r50_3.pth”