Skip to content

Official code for "Binary embedding based retrieval at Tencent"

License

Notifications You must be signed in to change notification settings

TencentARC/BEBR

Repository files navigation

BEBR

Approach

BEBR

Usage

First, install Pytorch 1.13.1 (or later) and torchvision, as well as some additional dependencies

$ conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia
$ git clone https://github.com/ganyk/BEBR.git
$ pip install -r requirements

Evaluation

Pre-computed hash features, recurrent binary features and float features are provided for evaluation.

# run evaluation on hash features
$ python tools/eval.py --image_feat dataset/hash/feat_image.npy --txt_feat dataset/hash/feat_txt.npy

# float features
$ python tools/eval.py --image_feat dataset/float_finetune/feat_image.npy --txt_feat dataset/float_finetune/feat_txt.npy

# recurrent binary features
$ python tools/eval.py --image_feat dataset/bebr/feat_image.npy --txt_feat dataset/bebr/feat_txt.npy

The results should be consistent with those in the paper:

Embedding Bits Recall@1 Recall@5 Recall@10
hash 1024 0.348 0.632 0.730
bebr 1024 0.360 0.646 0.751
float 16384 0.361 0.649 0.744

Train

Prepare data

  1. Download COCO datasets and uncompress them
$ mkdir dataset/coco & cd dataset/coco
$ wget http://images.cocodataset.org/zips/train2017.zip
$ wget http://images.cocodataset.org/zips/val2017.zip
$ wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip

unzip train2017.zip & unzip val2017.zip & unzip annotations_trainval2017.zip
  1. Use clip RN101 model to generate float features which will be used as input to train binary model
$ python tools/gen_clip_feat_train.py
$ python tools/gen_clip_feat_eval.py

Train binary model

# train hash model
$ sh train_local.sh configs/hash.yaml
# train bebr model
$ sh train_local.sh configs/bebr.yaml

finetune float feature

The training process of binary model can be seened as a finetune process. For fair comparison, we also finetune the clip features using MLPs whose input and output are both float features.

$ sh train_local.sh configs/float.yaml

About

Official code for "Binary embedding based retrieval at Tencent"

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published