UniT: Unified Knowledge Transfer for Any-shot Object Detection and Segmentation

This repository contains the code for the CVPR 2021 paper titled "UniT: Unified Knowledge Transfer for Any-shot Object Detection and Segmentation".

Requirements

To setup the environment with all the required dependencies, follow the steps detailed in INSTALL.md.

Prepare Dataset

To obtain the data and curate the splits for PASCAL VOC

python data/prepare_voc.py --DATA-ROOT "Path/to/Save/Location/"

Note: These splits are the same as the ones released by the authors of Few-shot Object Detection via Feature Reweighting.

Model Training

Download ImageNet pretrained models from this Google drive, and place them in the models/ folder in the root directory. The directory structure should look like this:

UniT
└── models
    └── resnet_50_MSRA_C4.pth
    └── resnet_101_MSRA_C4.pth
    ...
└── modeling
└── scripts
└── solver
...

Base Training

Training on VOC

python scripts/train_VOC.py --config-file "configs/VOC/VOC-RCNN-101-C4-split{num}.yaml" --num-gpus 4 --resume SOLVER.IMS_PER_BATCH 8 TEST.AUG.ENABLED False SOLVER.BASE_LR 0.02

Fine Tuning

Fine Tuning on VOC

python scripts/finetune_VOC.py --config-file "configs/VOC/FT/{num}_shot/VOC-RCNN-101-C4-split{num}-ft.yaml" --num-gpus 4 --resume OUTPUT_DIR "Path/for/Checkpointing" MODEL.WEIGHTS "Path/to/Base/Training/Model/Weights"

Evaluation

Evaluation on VOC Evaluation can be done at any stage using the --eval-only flag. For example, the model obtained after fine tuning can be evaluated as follows:

python scripts/finetune_VOC.py --config-file "configs/VOC/VOC-RCNN-101-C4-split{num}.yaml" --num-gpus 4 --eval-only --resume OUTPUT_DIR "Path/for/Checkpointing" MODEL.WEIGHTS "Path/to/Fine/Tune/Model/Weights"

Note: The default training/testing assumes 4 GPUs. It can be modified to suit other GPU configurations, but would require changing the learning rate and batch sizes accordingly. Please look at SOLVER.REFERENCE_WORLD_SIZE parameter in the detectron2 configurations for details on how this can be done automatically.

This repository is still being updated. Instructions on how to run the code on MS-COCO will be provided shortly.

ubc-vision / UniT