yblir / rt-detr

TODO

see details

• Training
• Evaluation
• Export onnx
• Upload source code
• Upload weight convert from paddle, see links
• Align training details with the paddle version
• Tuning rtdetr based on pretrained weights

Model Zoo

Model	Dataset	Input Size	APval	AP50val	#Params(M)	FPS	checkpoint
rtdetr_r18vd	COCO	640	46.4	63.7	20	217	url*
rtdetr_r34vd	COCO	640	48.9	66.8	31	161	url*
rtdetr_r50vd_m	COCO	640	51.3	69.5	36	145	url*
rtdetr_r50vd	COCO	640	53.1	71.2	42	108	url*
rtdetr_r101vd	COCO	640	54.3	72.8	76	74	url*
rtdetr_18vd	COCO+Objects365	640	49.0	66.5	20	217	url*
rtdetr_r50vd	COCO+Objects365	640	55.2	73.4	42	108	url*
rtdetr_r101vd	COCO+Objects365	640	56.2	74.5	76	74	url*

Notes

• COCO + Objects365 in the table means finetuned model on COCO using pretrained weights trained on Objects365.
• url* is the url of pretrained weights convert from paddle model for save energy. It may have slight differences between this table and paper

Quick start

Install

pip install -r requirements.txt

Data

• Download and extract COCO 2017 train and val images.

path/to/coco/
  annotations/  # annotation json files
  train2017/    # train images
  val2017/      # val images

• Modify config img_folder, ann_file

Training & Evaluation

• Training on a Single GPU:

# training on single-gpu
export CUDA_VISIBLE_DEVICES=0
python tools/train.py -c configs/rtdetr/rtdetr_r50vd_6x_coco.yml

• Training on Multiple GPUs:

# train on multi-gpu
export CUDA_VISIBLE_DEVICES=0,1,2,3
torchrun --nproc_per_node=4 tools/train.py -c configs/rtdetr/rtdetr_r50vd_6x_coco.yml

• Evaluation on Multiple GPUs:

# val on multi-gpu
export CUDA_VISIBLE_DEVICES=0,1,2,3
torchrun --nproc_per_node=4 tools/train.py -c configs/rtdetr/rtdetr_r50vd_6x_coco.yml -r path/to/checkpoint --test-only

Export

python tools/export_onnx.py -c configs/rtdetr/rtdetr_r18vd_6x_coco.yml -r path/to/checkpoint --check

Train custom data

1. set remap_mscoco_category: False. This variable only works for ms-coco dataset. If you want to use remap_mscoco_category logic on your dataset, please modify variable mscoco_category2name based on your dataset.

2. add -t path/to/checkpoint (optinal) to tuning rtdetr based on pretrained checkpoint. see training script details.