README.md
March 8, 2023 ยท View on GitHub
KS-DETR: Knowledge Sharing in Attention Learning for Detection Transformer
We release our code for our submitted manuscript KS-DETR: Knowledge Sharing in Attention Learning for Detection Transformer.
Main results and Pretrained Models
Here we provide the pretrained KS-DETR weights based on detrex.
| Name | Backbone | Pretrain | Epochs | box AP |
download |
|---|---|---|---|---|---|
| KS-DAB-DETR-R50 | R-50 | IN1k | 50 | 43.9 | model |
| KS-DAB-DETR-R101 | R-101 | IN1k | 50 | 45.3 | model |
| KS-DAB-DETR-Swin-T | Swin-T | IN1k | 50 | 47.1 | model |
| KS-Conditional-DETR-R50 | R-50 | IN1k | 50 | 45.3 | model |
| KS-Conditional-DETR-R101 | R-101 | IN1k | 50 | 47.1 | model |
| KS-DN-DETR-R50 | R-50 | IN1k | 50 | 45.2 | model |
| KS-DN-DETR-R101 | R-101 | IN1k | 50 | 46.5 | model |
| KS-Deformable-DETR-R50 | R-50 | IN1k | 12 | 36.4 | model |
| KS-Deformable-DETR-R101 | R-101 | IN1k | 12 | 38.4 | model |
| KS-DN-Deformable-DETR-R50 | R-50 | IN1k | 12 | 46.5 | model |
| KS-Deformable-DETR-R50 | R-50 | IN1k | 50 | 44.8 | model |
| KS-Deformable-DETR-R101 | R-101 | IN1k | 50 | 46.0 | model |
Installation
conda create -n ksdetr python=3.8 -y
conda activate ksdetr
git clone https://github.com/edocanonymous/KS-DETR
cd KS-DETR
python -m pip install -e detectron2
pip install -e .
Training
To train the models with R101 backbone, the pretrained IN1k weights should be available at location output/weights/R-101.pkl.
We can follow https://github.com/facebookresearch/detectron2/blob/main/tools/convert-torchvision-to-d2.py
to convert https://download.pytorch.org/models/resnet101-5d3b4d8f.pth
to torchvision format and obtain R-101.pkl by
wget https://download.pytorch.org/models/resnet101-5d3b4d8f.pth -O output/r101.pth
python ./detectron2/tools/convert-torchvision-to-d2.py output/r101.pth output/weights/R-101.pkl
We provide our converted R-101.pkl file here.
All configs can be trained with:
cd detrex
python tools/train_net.py --config-file projects/dab_detr/configs/path/to/config.py --num-gpus 8
To train KS-DAB-DETR-R50, KS-DAB-DETR-R101, and KS-DAB-DETR-Swin-T,
python tools/train_net.py --config-file projects/ks_detr/configs/ks_dab_detr/ks_dab_detr_r50_50ep_smlp_qkv_triple_attn.py --num-gpus 8
python tools/train_net.py --config-file projects/ks_detr/configs/ks_dab_detr/ks_dab_detr_r101_50ep_smlp_qkv_triple_attn.py --num-gpus 8
python tools/train_net.py --config-file projects/ks_detr/configs/ks_dab_detr/ks_dab_detr_swin_tiny_50ep_smlp_qkv_triple_attn.py --num-gpus 8
Evaluation
Model evaluation can be done as follows:
cd detrex
python tools/train_net.py --config-file projects/dab_detr/configs/path/to/config.py --eval-only train.init_checkpoint=/path/to/model_checkpoint
License
This project is released under the Apache 2.0 license.
Acknowledgement
-
Our code is built on detrex, which is an open-source toolbox for Transformer-based detection algorithms created by researchers of IDEACVR.
-
detrex is built based on Detectron2 and part of its module design is borrowed from MMDetection, DETR, and Deformable-DETR.