Code for CVPR 2018 Paper "HashGAN: Deep Learning to Hash with Pair Conditional Wasserstein GAN".
- Python3, NumPy, TensorFlow-gpu, SciPy, Matplotlib, OpenCV, easydict, yacs, tqdm
- A recent NVIDIA GPU
We provide a environment.yaml
for you and you can simplely use conda env create -f environment.yml
to create the environment.
Or you can create the environment from scratch:
conda create --no-default-packages -n HashGAN python=3.6 && source activate HashGAN
conda install -y numpy scipy matplotlib tensorflow-gpu opencv
pip install easydict yacs tqdm pillow
In data_list/
folder, we give three examples to show how to prepare image training data. If you want to add other datasets as the input, you need to prepare train.txt
, test.txt
, database.txt
and database_nolabel.txt
as CIFAR-10 dataset.
You can download the whole cifar10 dataset including the images and data list from here, and unzip it to data/cifar10
folder.
If you need run on NUSWIDE_81 and COCO, we recommend you to follow here to prepare NUSWIDE_81 and COCO images.
The imagenet pretrained Alexnet model can be downloaded here. You can download the pretrained Generator models in the release page and modify config file to use the pretrained models.
The training process can be divided into two step:
- Training a image generator.
- Fintune Alexnet using original labeled images and generated images.
In config
folder, we provide some examples to prepare yaml configuration.
config
├── cifar_evaluation.yaml
├── cifar_step_1.yaml
├── cifar_step_2.yaml
└── nuswide_step_1.yaml
You can run the model using command like the following:
python main.py --cfg config/cifar_step_1.yaml --gpus 0
python main.py --cfg config/cifar_step_2.yaml --gpus 0
You can use tensorboard to monitor the training process such as losses and Mean Average Precision.
If you use this code for your research, please consider citing:
@inproceedings{cao2018hashgan,
title={HashGAN: Deep Learning to Hash with Pair Conditional Wasserstein GAN},
author={Cao, Yue and Liu, Bin and Long, Mingsheng and Wang, Jianmin and KLiss, MOE},
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
pages={1287--1296},
year={2018}
}
If you have any problem about our code, feel free to contact
or describe your problem in Issues.