Home

Awesome

KakaoBrain pytorch pytorch

Sparse DETR (ICLR'22)

By Byungseok Roh*, Jaewoong Shin*, Wuhyun Shin*, and Saehoon Kim at Kakao Brain. (*: Equal contribution)

Introduction

TL; DR. Sparse DETR is an efficient end-to-end object detector that sparsifies encoder tokens by using the learnable DAM(Decoder Attention Map) predictor. It achieves better performance than Deformable DETR even with only 10% encoder queries on the COCO dataset.

<p align="center"> <img src="./figs/dam_creation.png" height=350> </p>

Abstract. DETR is the first end-to-end object detector using a transformer encoder-decoder architecture and demonstrates competitive performance but low computational efficiency on high resolution feature maps. The subsequent work, Deformable DETR, enhances the efficiency of DETR by replacing dense attention with deformable attention, which achieves 10x faster convergence and improved performance. Deformable DETR uses the multiscale feature to ameliorate performance, however, the number of encoder tokens increases by 20x compared to DETR, and the computation cost of the encoder attention remains a bottleneck. In our preliminary experiment, we observe that the detection performance hardly deteriorates even if only a part of the encoder token is updated. Inspired by this observation, we propose Sparse DETR that selectively updates only the tokens expected to be referenced by the decoder, thus help the model effectively detect objects. In addition, we show that applying an auxiliary detection loss on the selected tokens in the encoder improves the performance while minimizing computational overhead. We validate that Sparse DETR achieves better performance than Deformable DETR even with only 10% encoder tokens on the COCO dataset. Albeit only the encoder tokens are sparsified, the total computation cost decreases by 38% and the frames per second (FPS) increases by 42% compared to Deformable DETR.

Installation

Requirements

We have tested the code on the following environments:

Run the following command to install dependencies:

pip install -r requirements.txt

Compiling CUDA operators

cd ./models/ops
sh ./make.sh
# unit test (should see all checking is True)
python test.py

Usage

Dataset preparation

Please download COCO 2017 dataset and organize them as follows:

code_root/
└── data/
    └── coco/
        ├── train2017/
        ├── val2017/
        └── annotations/
        	├── instances_train2017.json
        	└── instances_val2017.json

Training

Training on a single node

For example, the command for training Sparse DETR with the keeping ratio of 10% on 8 GPUs is as follows:

$ GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 8 ./configs/swint_sparse_detr_rho_0.1.sh

Training on multiple nodes

For example, the command Sparse DETR with the keeping ratio of 10% on 2 nodes of each with 8 GPUs is as follows:

On node 1:

$ MASTER_ADDR=<IP address of node 1> NODE_RANK=0 GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 16 ./configs/swint_sparse_detr_rho_0.1.sh

On node 2:

$ MASTER_ADDR=<IP address of node 2> NODE_RANK=1 GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 16 ./configs/swint_sparse_detr_rho_0.1.sh

Direct argument control

# Deformable DETR (with bounding-box-refinement and two-stage argument, if wanted)
$ GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 8 python main.py --with_box_refine --two_stage
# Efficient DETR (with the class-specific head as describe in their paper)
$ GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 8 python main.py --with_box_refine --two_stage --eff_query_init --eff_specific_head
# Sparse DETR (with the keeping ratio of 10% and encoder auxiliary loss)
$ GPUS_PER_NODE=8 ./tools/run_dist_launch.sh 8 python main.py --with_box_refine --two_stage --eff_query_init --eff_specific_head --rho 0.1 --use_enc_aux_loss

Some tips to speed-up training

Evaluation

You can get the pre-trained model of Sparse DETR (the link is in "Main Results" session), then run the following command to evaluate it on COCO 2017 validation set:

# Note that you should run the command with the corresponding configuration.
$ ./configs/swint_sparse_detr_rho_0.1.sh --resume <path to pre-trained model> --eval

You can also run distributed evaluation by using ./tools/run_dist_launch.sh.

Main Results

The tables below demonstrate the detection performance of Sparse DETR on the COCO 2017 validation set when using different backbones.

ResNet-50 backbone

MethodEpochsρTop-k & BBRAP#Params(M)GFLOPsB4FPSDownload
Faster R-CNN + FPN109N/A42.042M180G26
DETR50N/A35.041M86G28
DETR500N/A42.041M86G28
DETR-DC5500N/A43.341M187G12
PnP-DETR50033%41.1
50050%41.8
PnP-DETR-DC550033%42.7
50050%43.1
Deformable-DETR50N/A43.939.8M172.9G19.1
50N/Ao46.040.8M177.3G18.2
Sparse-DETR5010%o45.340.9M105.4G26.5link
5020%o45.640.9M112.9G24.8link
5030%o46.040.9M120.5G23.2link
5040%o46.240.9M128.0G21.8link
5050%o46.340.9M135.6G20.5link

Swin-T backbone

MethodEpochsρTop-k & BBRAP#Params(M)GFLOPsB4FPSDownload
DETR50N/A35.945.0M91.6G26.8
DETR500N/A45.445.0M91.6G26.8
Deformable-DETR50N/A45.740.3M180.4G15.9
50N/Ao48.041.3M184.8G15.4
Sparse-DETR5010%o48.241.4M113.4G21.2link
5020%o48.841.4M121.0G20link
5030%o49.141.4M128.5G18.9link
5040%o49.241.4M136.1G18link
5050%o49.341.4M143.7G17.2link

Initializing ResNet-50 backbone with SCRL

The performance of Sparse DETR can be further improved when the backbone network is initialized with the SCRL(Spatially Consistent Representation Learning) that aims to learn dense representations in a self-supervised way, compared to the default initialization with the ImageNet pre-trained one, denoted as IN-sup in the table below.

MethodρAP(IN-sup)AP(SCRL)AP(gain)Download
Sparse DETR10%45.346.9+1.6link
20%45.647.2+1.7link
30%46.047.4+1.4link
40%46.247.7+1.5link
50%46.347.9+1.6link

Citation

If you find Sparse DETR useful in your research, please consider citing:

@inproceedings{roh2022sparse,
  title={Sparse DETR: Efficient End-to-End Object Detection with Learnable Sparsity},
  author={Roh, Byungseok and Shin, JaeWoong and Shin, Wuhyun and Kim, Saehoon},
  booktitle={ICLR},
  year={2022}
}

License

This project is released under the Apache 2.0 license. Copyright 2021 Kakao Brain Corp. All Rights Reserved.