iCVTEAM / TDRG

Transformer-based Dual Relation Graph for Multi-label Image Recognition. ICCV 2021

TDRG

Pytorch implementation of Transformer-based Dual Relation Graph for Multi-label Image Recognition. ICCV 2021

Prerequisites

Python 3.6+

Pytorch 1.6

CUDA 10.1

Tesla V100 × 2

Datasets

• MS-COCO: train val annotations
• VOC 2007: trainval test test_anno

Model

• MS-COCO: the name of parameters in our original network is different from the public one, hence if you want to test the performance of TDRG on MS-COCO, please download the checkpoint into checkpoint/COCO2014 folder and replace the function load_checkpoint with load_origin_checkpoint in trainer.py.

Train

CUDA_VISIBLE_DEVICES=0,1 python main.py  --data COCO2014 --data_root_dir $DATA_PATH$ --save_dir $SAVE_PATH$ --i 448  --lr 0.03 -b 64

Test

python main.py  --data COCO2014 --data_root_dir $DATA_PATH$ --save_dir $SAVE_PATH$ --i 448  --lr 0.03 -b 64 -e --resume checkpoint/COCO2014/checkpoint_COCO.pth

Visualization

Citation

• If you find this work is helpful, please cite our paper

@InProceedings{Zhao2021TDRG,
    author    = {Zhao, Jiawei and Yan, Ke and Zhao, Yifan and Guo, Xiaowei and Huang, Feiyue and Li, Jia},
    title     = {Transformer-Based Dual Relation Graph for Multi-Label Image Recognition},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {163-172}
}