This repository is written by Chenhan Jiang and Hang Xu interned at SenseTime
Code for reproducing the results in the following paper, and large part code is reference from jwyang/faster-rcnn.pytorch:
Hybrid Knowledge Routed Modules for Large-scale Object Detection
Chenhan Jiang, Hang Xu, Xiaodan Liang, Liang Lin
Neural Information Processing Systems (NIPS), 2018
We also provide Supplementary Materials download address.
Clone the repo:
https://github.com/chanyn/HKRM.git
python packages
This project can not support pytorch 0.4, higher version will not recur results.
Torchvision >= 0.2.0
cython
pyyaml
easydict
opencv-python
matplotlib
numpy
scipy
tensorboardX
You can install above package using pip
:
pip install Cython easydict matplotlib opencv-python pyyaml scipy
CUDA 8.0
gcc >= 4.9
Compile the CUDA dependencies:
cd {repo_root}/lib
sh make.sh
It will compile all the modules you need, including NMS, ROI_Pooing, ROI_Crop and ROI_Align.
Create a data folder under the repo,
cd {repo_root}
mkdir data
Prior Knowledge: We instantiate attribute and relation as explicit knowledge from Visual Genome. We provide the link to download VG graph and other graph. We also provide the frequency statistic from VG and procession code in lib/dataset/tool/compute_prior.py
. So You can transfer prior knowledge to other datasets.
ADE: We provide ADE20K as an example.
mkdir -p data/ADE
cd data/ADE
wget -v http://groups.csail.mit.edu/vision/datasets/ADE20K/ADE20K_2016_07_26.zip
tar -xzvf ADE20K_2016_07_26.zip
mv ADE20K_2016_07_26/* ./
rmdir ADE20K_2016_07_26
# then get the list of overlap with VG and train/test split
wget -v https://drive.google.com/file/d/1rbN0MZM4ome7gk4hRAziJ60IivkMYWbm/view?usp=sharing
tar -xzvf ADE_split.tar.gz
rm -vf ADE_split.tar.gz
cd ../..
Visual Genome: Download the VG images and annotations from Visual Genome. We use synset as label rather than name, and choose top 1000 and 3000 frequent classes. Our JSON datasets can be download from:
Other Datasets: We also implement our module on MSCOCO and PASCAL VOC. The classes of these datasets overlap with Visual Genome top 3000 frequent categories. So we get corresponding graph from VG3000. You can directly download MSCOCO and PACAL VOC graphs, or you can follow the transfer example in lib/tools/collect_coco_voc_gtgraph.py.
We used ResNet101 pretrained model on ImageNet in our experiments. Download it and put it into the data/pretrained_model/.
For example, to train the baseline with res101 on VG, simply run:
CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_baseline.py \
--dataset vg --bs $BATCH_SIZE --nw $WORKER_NUMBER \
--log_dir $LOG_DIR --save_dir $WHERE_YOU_WANT
where 'bs' is the batch size with default 2, 'log_dir' is the location to save your tensorboard of train, 'save_dir' is your path to save model.
Change dataset to 'ade' or 'coco' if you want to train on ADE or COCO. If you have multi-GPUs, you should add '--mGPUs'.
To train our HKRM with res101 on VG:
CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_HKRM.py \
--dataset vg --bs $BATCH_SIZE --nw $WORKER_NUMBER \
--log_dir $LOG_DIR --save_dir $WHERE_YOU_WANT \
--init --net HKRM --attr_size 256 --rela_size 256 --spat_size 256
where 'net' you can choose in ('HKRM', 'Attribute', 'Relation', 'Spatial'), and you should put pretrained baseline model into 'save_dir'.
If you want to evaluate the detection performance of a pre-trained model on VG test set, simply run:
python test_net.py --dataset vg --net HKRM \
--load_dir $YOUR_SAVE_DIR \
--checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT
Specify the specific model session, chechepoch and checkpoint, e.g., SESSION=325, EPOCH=12, CHECKPOINT=21985. And we provide the final model that you can load from trained_model_hkrm.
@inproceedings{jiangNIPS18hybrid,
Author = {Chenhan Jiang and Hang Xu and Xiaodan Liang and Liang Lin},
Title = {Hybrid Knowledge Routed Modules for Large-scale Object Detection},
Booktitle = {Advances in Neural Information Processing Systems ({NIPS})},
Year = {2018}
}
If you have any questions about this repo, please feel free to contact jchcyan@gmail.com.