Few-Shot Object Detection via Association and DIscrimination

Code release of our NeurIPS 2021 paper: Few-Shot Object Detection via Association and DIscrimination.

Bibtex

@inproceedings{cao2021few,
  title={Few-Shot Object Detection via Association and DIscrimination},
  author={Cao, Yuhang and Wang, Jiaqi and Jin, Ying and Wu, Tong and Chen, Kai and Liu, Ziwei and Lin, Dahua},
  booktitle={Thirty-Fifth Conference on Neural Information Processing Systems},
  year={2021}
}

Arxiv: https://arxiv.org/abs/2111.11656

Install dependencies

• Create a new environment: conda create -n fadi python=3.8 -y
• Active the newly created environment: conda activate fadi
• Install PyTorch and torchvision: conda install pytorch=1.7 torchvision cudatoolkit=10.2 -c pytorch -y
• Install MMDetection: pip install mmdet==2.11.0
• Install MMCV: pip install mmcv==1.2.5
• Install MMCV-Full: pip install mmcv-full==1.2.5 -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.7.0/index.html

Note:

• Only tested on MMDet==2.11.0, MMCV==1.2.5, it may not be consistent with other versions.
• The above instructions use CUDA 10.2, make sure you install the correct PyTorch, Torchvision and MMCV-Full that are consistent with your CUDA version.

Prepare dataset

We follow exact the same split with TFA, please download the dataset and split files as follows:

• Download PASCAL VOC
• Download split files

Create a directory data in the root directory, and the expected structure for data directory:

data/
    VOCdevkit
    few_shot_voc_split

Training & Testing

Base Training

FADI share the same base training stage with TFA, we directly convert the corresponding checkpoints from TFA in Detectron2 format to MMDetection format, please download the base training checkpoints following the table.

Create a directory models in the root directory, and the expected structure for models directory:

models/
    voc_split1_base.pth
    voc_split1_base_surgery.pth
    voc_split2_base.pth
    voc_split2_base_surgery.pth
    voc_split3_base.pth
    voc_split3_base_surgery.pth

Few-Shot Fine-tuning

FADI divides the few-shot fine-tuning stage into two steps, ie, association and discrimination,

Suppose we want to train a model for Pascal VOC split1, shot1 with 8 GPUs

1. Step 1: Association.

Getting the assigning scheme of the split:

python tools/associate.py 1

Aligning the feature distribution of the associated base and novel classes:

./tools/dist_train.sh configs/voc_split1/fadi_split1_shot1_association.py 8

2. Step 2: Discrimination

Building a discriminate feature space for novel classes with disentangling and set-specialized margin loss:

./tools/dist_train.sh configs/voc_split1/fadi_split1_shot1_discrimination.py 8

Holistically Training:

We also provide you a script tools/fadi_finetune.sh to holistically train a model for a specific split/shot by running:

./tools/fadi_finetune.sh 1 1

Evaluation

To evaluate the trained models, run

./tools/dist_test.sh configs/voc_split1/fadi_split1_shot1_discrimination.py [checkpoint] 8 --eval mAP --out res.pkl

Model Zoo

Pascal VOC split 1

Pascal VOC split 2

Pascal VOC split 3