GAN Compression

project | paper | videos | slides

[NEW!] GAN Compression is accepted by T-PAMI! We released our T-PAMI version in the arXiv v4!

[NEW!] We release the codes of our interactive demo and include the TVM tuned model. It achieves 8FPS on Jetson Nano GPU now!

[NEW!] Add support to the MUNIT, a multimodal unsupervised image-to-image translation approach! Please follow the test commands to test the pre-trained models and the tutorial to train your own models!

We introduce GAN Compression, a general-purpose method for compressing conditional GANs. Our method reduces the computation of widely-used conditional GAN models, including pix2pix, CycleGAN, MUNIT, and GauGAN, by 9-29x while preserving the visual fidelity. Our method is effective for a wide range of generator architectures, learning objectives, and both paired and unpaired settings.

GAN Compression: Efficient Architectures for Interactive Conditional GANs
Muyang Li, Ji Lin, Yaoyao Ding, Zhijian Liu, Jun-Yan Zhu, and Song Han
MIT, Adobe Research, SJTU
In CVPR 2020.

Demos

Overview

GAN Compression framework: ① Given a pre-trained teacher generator G', we distill a smaller “once-for-all” student generator G that contains all possible channel numbers through weight sharing. We choose different channel numbers for the student generator G at each training step. ② We then extract many sub-generators from the “once-for-all” generator and evaluate their performance. No retraining is needed, which is the advantage of the “once-for-all” generator. ③ Finally, we choose the best sub-generator given the compression ratio target and performance target (FID or mIoU) using either brute-force or evolutionary search method. Optionally, we perform additional fine-tuning, and obtain the final compressed model.

Performance

GAN Compression reduces the computation of pix2pix, cycleGAN and GauGAN by 9-21x, and model size by 4.6-33x.

Colab Notebook

PyTorch Colab notebook: CycleGAN and pix2pix.

Prerequisites

Linux
Python 3
CPU or NVIDIA GPU + CUDA CuDNN

Getting Started

Installation

Clone this repo:

git clone git@github.com:mit-han-lab/gan-compression.git
cd gan-compression

Install PyTorch 1.4 and other dependencies (e.g., torchvision).
- For pip users, please type the command pip install -r requirements.txt.
- For Conda users, we provide an installation script scripts/conda_deps.sh. Alternatively, you can create a new Conda environment using conda env create -f environment.yml.

CycleGAN

Setup

Download the CycleGAN dataset (e.g., horse2zebra).

bash datasets/download_cyclegan_dataset.sh horse2zebra

Get the statistical information for the ground-truth images for your dataset to compute FID. We provide pre-prepared real statistic information for several datasets. For example,
```
bash datasets/download_real_stat.sh horse2zebra A
bash datasets/download_real_stat.sh horse2zebra B
```

Apply a Pre-trained Model

Download the pre-trained models.

python scripts/download_model.py --model cycle_gan --task horse2zebra --stage full
python scripts/download_model.py --model cycle_gan --task horse2zebra --stage compressed

Test the original full model.

bash scripts/cycle_gan/horse2zebra/test_full.sh

Test the compressed model.

bash scripts/cycle_gan/horse2zebra/test_compressed.sh

Measure the latency of the two models.

bash scripts/cycle_gan/horse2zebra/latency_full.sh
bash scripts/cycle_gan/horse2zebra/latency_compressed.sh

There may be a little differences between the results of above models and those of the paper since we retrained the models. We also release the compressed models of the paper. If there are such inconsistencies, you could try the following commands to test our paper models:
```
python scripts/download_model.py --model cycle_gan --task horse2zebra --stage legacy
bash scripts/cycle_gan/horse2zebra/test_legacy.sh
bash scripts/cycle_gan/horse2zebra/latency_legacy.sh
```

Pix2pix

Setup

Download the pix2pix dataset (e.g., edges2shoes).

bash datasets/download_pix2pix_dataset.sh edges2shoes-r

Get the statistical information for the ground-truth images for your dataset to compute FID. We provide pre-prepared real statistics for several datasets. For example,
```
bash datasets/download_real_stat.sh edges2shoes-r B
bash datasets/download_real_stat.sh edges2shoes-r subtrain_B
```

Apply a Pre-trained Model

Download the pre-trained models.

python scripts/download_model.py --model pix2pix --task edges2shoes-r --stage full
python scripts/download_model.py --model pix2pix --task edges2shoes-r --stage compressed

Test the original full model.

bash scripts/pix2pix/edges2shoes-r/test_full.sh

Test the compressed model.

bash scripts/pix2pix/edges2shoes-r/test_compressed.sh

Measure the latency of the two models.

bash scripts/pix2pix/edges2shoes-r/latency_full.sh
bash scripts/pix2pix/edges2shoes-r/latency_compressed.sh

There may be a little differences between the results of above models and those of the paper since we retrained the models. We also release the compressed models of the paper. If there are such inconsistencies, you could try the following commands to test our paper models:
```
python scripts/download_model.py --model pix2pix --task edges2shoes-r --stage legacy
bash scripts/pix2pix/edges2shoes-r/test_legacy.sh
bash scripts/pix2pix/edges2shoes-r/latency_legacy.sh
```

GauGAN

Setup

Prepare the cityscapes dataset. Check here for preparing the cityscapes dataset.
Get the statistical information for the ground-truth images for your dataset to compute FID. We provide pre-prepared real statistics for several datasets. For example,
```
bash datasets/download_real_stat.sh cityscapes A
```

Apply a Pre-trained Model

Download the pre-trained models.

python scripts/download_model.py --model gaugan --task cityscapes --stage full
python scripts/download_model.py --model gaugan --task cityscapes --stage compressed

Test the original full model.

bash scripts/gaugan/cityscapes/test_full.sh

Test the compressed model.

bash scripts/gaugan/cityscapes/test_compressed.sh

Measure the latency of the two models.

bash scripts/gaugan/cityscapes/latency_full.sh
bash scripts/gaugan/cityscapes/latency_compressed.sh

There may be a little differences between the results of above models and those of the paper since we retrained the models. We also release the compressed models of the paper. If there are such inconsistencies, you could try the following commands to test our paper models:
```
python scripts/download_model.py --model gaugan --task cityscapes --stage legacy
bash scripts/gaugan/cityscapes/test_legacy.sh
bash scripts/gaugan/cityscapes/latency_legacy.sh
```

MUNIT

Setup

Prepare the dataset (e.g., edges2shoes-r).

bash datasets/download_pix2pix_dataset.sh edges2shoes-r
python datasets/separate_A_and_B.py --input_dir database/edges2shoes-r --output_dir database/edges2shoes-r-unaligned
python datasets/separate_A_and_B.py --input_dir database/edges2shoes-r --output_dir database/edges2shoes-r-unaligned --phase val

Get the statistical information for the ground-truth images for your dataset to compute FID. We provide pre-prepared real statistics for several datasets. For example,
```
bash datasets/download_real_stat.sh edges2shoes-r B
bash datasets/download_real_stat.sh edges2shoes-r-unaligned subtrain_B
```

Apply a Pretrained Model

Download the pre-trained models.

python scripts/download_model.py --model gaugan --task cityscapes --stage full
python scripts/download_model.py --model gaugan --task cityscapes --stage compressed

Test the original full model.

bash scripts/munit/edges2shoes-r_fast/test_full.sh

Test the compressed model.

bash scripts/munit/edges2shoes-r_fast/test_compressed.sh

Measure the latency of the two models.

bash scripts/munit/edges2shoes-r_fast/latency_full.sh
bash scripts/munit/edges2shoes-r_fast/latency_compressed.sh

Cityscapes Dataset

For the Cityscapes dataset, we cannot provide it due to license issue. Please download the dataset from https://cityscapes-dataset.com and use the script prepare_cityscapes_dataset.py to preprocess it. You need to download gtFine_trainvaltest.zip and leftImg8bit_trainvaltest.zip and unzip them in the same folder. For example, you may put gtFine and leftImg8bit in database/cityscapes-origin. You need to prepare the dataset with the following commands:

python datasets/get_trainIds.py database/cityscapes-origin/gtFine/
python datasets/prepare_cityscapes_dataset.py \
--gtFine_dir database/cityscapes-origin/gtFine \
--leftImg8bit_dir database/cityscapes-origin/leftImg8bit \
--output_dir database/cityscapes \
--train_table_path datasets/train_table.txt \
--val_table_path datasets/val_table.txt

You will get a preprocessed dataset in database/cityscapes and a mapping table (used to compute mIoU) in dataset/table.txt.

To support mIoU computation, you need to download a pre-trained DRN model drn-d-105_ms_cityscapes.pth from http://go.yf.io/drn-cityscapes-models. By default, we put the drn model in the root directory of the repo. Then you can test our compressed models on cityscapes after you have downloaded our models.

COCO-Stuff Dataset

We follow the same COCO-Stuff dataset preparation as NVlabs/spade. Specifically, you need to download train2017.zip, val2017.zip, stuffthingmaps_trainval2017.zip, and annotations_trainval2017.zip from nightrome/cocostuff. The images, labels, and instance maps should be arranged in the same directory structure as in datasets/coco_stuff. In particular, we used an instance map that combines both the boundaries of "things instance map" and "stuff label map". To do this, we used a simple script datasets/coco_generate_instance_map.py.

To support mIoU computation, you need to download a pre-trained DeeplabV2 model deeplabv2_resnet101_msc-cocostuff164k-100000.pth and also put it in the root directory of the repo.

Performance of Released Models

Here we show the performance of all our released models:

Model	Dataset	Method	#Parameters	MACs	Metric
Model	Dataset	Method	#Parameters	MACs	FID	mIoU
CycleGAN	horse→zebra	Original	11.4M	56.8G	65.75	--
		GAN Compression (Paper)	0.342M	2.67G	65.33	--
		GAN Compression (Retrained)	0.357M	2.55G	65.12	--
		Fast GAN Compression	0.355M	2.64G	65.19	--
Pix2pix	edges→shoes	Original	11.4M	56.8G	24.12	--
		GAN Compression (Paper)	0.700M	4.81G	26.60	--
		GAN Compression (Retrained)	0.822M	4.99G	26.70	--
		Fast GAN Compression	0.703M	4.83G	25.76	--
	Cityscapes	Original	11.4M	56.8G	--	42.06
		GAN Compression (Paper)	0.707M	5.66G	--	40.77
		GAN Compression (Retrained)	0.781M	5.59G	--	38.63
		Fast GAN Compression	0.867M	5.61G	--	41.71
	map→arial photo	Original	11.4M	56.8G	47.91	--
		GAN Compression	0.746M	4.68G	48.02	--
		Fast GAN Compression	0.708M	4.53G	48.67	--
GauGAN	Cityscapes	Original	93.0M	281G	57.60	61.04
		GAN Compression (Paper)	20.4M	31.7G	55.19	61.22
		GAN Compression (Retrained)	21.0M	31.2G	56.43	60.29
		Fast GAN Compression	20.2M	31.3G	56.25	61.17
	COCO-Stuff	Original	97.5M	191G	21.38	38.78
	COCO-Stuff	Fast GAN Compression	26.0M	35.5G	25.06	35.05
MUNIT	edges→shoes	Original	15.0M	77.3G	30.13	--
MUNIT	edges→shoes	Fast GAN Compression	1.10M	2.63G	30.53	--

Training

Please refer to the tutorial of Fast GAN Compression and GAN Compression on how to train models on our datasets and your own.

FID Computation

To compute the FID score, you need to get some statistical information from the groud-truth images of your dataset. We provide a script get_real_stat.py to extract statistical information. For example, for the edges2shoes dataset, you could run the following command:

python get_real_stat.py \
--dataroot database/edges2shoes-r \
--output_path real_stat/edges2shoes-r_B.npz \
--direction AtoB

For paired image-to-image translation (pix2pix and GauGAN), we calculate the FID between generated test images to real test images. For unpaired image-to-image translation (CycleGAN), we calculate the FID between generated test images to real training+test images. This allows us to use more images for a stable FID evaluation, as done in previous unconditional GANs research. The difference of the two protocols is small. The FID of our compressed CycleGAN model increases by 4 when using real test images instead of real training+test images.

Code Structure

To help users better understand and use our code, we briefly overview the functionality and implementation of each package and each module.

Citation

If you use this code for your research, please cite our paper.

@inproceedings{li2020gan,
  title={GAN Compression: Efficient Architectures for Interactive Conditional GANs},
  author={Li, Muyang and Lin, Ji and Ding, Yaoyao and Liu, Zhijian and Zhu, Jun-Yan and Han, Song},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2020}
}

Acknowledgements

Our code is developed based on pytorch-CycleGAN-and-pix2pix, SPADE, and MUNIT.

We also thank pytorch-fid for FID computation, drn for cityscapes mIoU computation and deeplabv2 for Coco-Stuff mIoU computation.

Name		Name	Last commit message	Last commit date
Latest commit History 128 Commits
configs		configs
data		data
datasets		datasets
distillers		distillers
docs		docs
imgs		imgs
interactive_demo		interactive_demo
metric		metric
models		models
options		options
scripts		scripts
supernets		supernets
utils		utils
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
assemble.py		assemble.py
cycle_gan.ipynb		cycle_gan.ipynb
distill.py		distill.py
environment.yaml		environment.yaml
evolution_search.py		evolution_search.py
export.py		export.py
get_real_stat.py		get_real_stat.py
get_test_opt.py		get_test_opt.py
latency.py		latency.py
merge.py		merge.py
pix2pix.ipynb		pix2pix.ipynb
remove_spectral_norm.py		remove_spectral_norm.py
requirements.txt		requirements.txt
search.py		search.py
select_arch.py		select_arch.py
test.py		test.py
tox.ini		tox.ini
train.py		train.py
train_supernet.py		train_supernet.py
trainer.py		trainer.py

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Repository files navigation

GAN Compression

project | paper | videos | slides

Demos

Overview

Performance

Colab Notebook

Prerequisites

Getting Started

Installation

CycleGAN

Setup

Apply a Pre-trained Model

Pix2pix

Setup

Apply a Pre-trained Model

GauGAN

Setup

Apply a Pre-trained Model

MUNIT

Setup

Apply a Pretrained Model

Cityscapes Dataset

COCO-Stuff Dataset

Performance of Released Models

Training

FID Computation

Code Structure

Citation

Acknowledgements

About

Releases 7

Packages

Contributors 4

Languages

License

mit-han-lab/gan-compression

Folders and files

Latest commit

History

Repository files navigation

GAN Compression

project | paper | videos | slides

Demos

Overview

Performance

Colab Notebook

Prerequisites

Getting Started

Installation

CycleGAN

Setup

Apply a Pre-trained Model

Pix2pix

Setup

Apply a Pre-trained Model

GauGAN

Setup

Apply a Pre-trained Model

MUNIT

Setup

Apply a Pretrained Model

Cityscapes Dataset

COCO-Stuff Dataset

Performance of Released Models

Training

FID Computation

Code Structure

Citation

Acknowledgements

About

Topics

Resources

License

Stars

Watchers

Forks

Releases 7

Packages 0

Contributors 4

Languages

Packages