Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/mit-han-lab/gan-compression

[CVPR 2020] GAN Compression: Efficient Architectures for Interactive Conditional GANs
https://github.com/mit-han-lab/gan-compression

compression conditional-gans cyclegan gans gaugan image-to-image-translation pix2pix pytorch

Last synced: 5 days ago
JSON representation

[CVPR 2020] GAN Compression: Efficient Architectures for Interactive Conditional GANs

Host: GitHub
URL: https://github.com/mit-han-lab/gan-compression
Owner: mit-han-lab
License: other
Created: 2020-03-05T00:27:13.000Z (almost 5 years ago)
Default Branch: master
Last Pushed: 2024-06-05T21:59:17.000Z (7 months ago)
Last Synced: 2024-12-13T11:11:48.839Z (12 days ago)
Topics: compression, conditional-gans, cyclegan, gans, gaugan, image-to-image-translation, pix2pix, pytorch
Language: Python
Homepage:
Size: 49.1 MB
Stars: 1,105
Watchers: 28
Forks: 150
Open Issues: 9
Metadata Files:
- Readme: README.md
- License: LICENSE

Awesome Lists containing this project

README

# GAN Compression
### [project](https://hanlab.mit.edu/projects/gancompression/) | [paper](https://arxiv.org/abs/2003.08936) | [videos](https://www.youtube.com/playlist?list=PL80kAHvQbh-r5R8UmXhQK1ndqRvPNw_ex) | [slides](https://hanlab.mit.edu/projects/gancompression/resources/546-slides.pdf)

**[NEW!]** GAN Compression is accepted by T-PAMI! We released our T-PAMI version in the [arXiv v4](https://arxiv.org/abs/2003.08936)!

**[NEW!]** We release the codes of our [interactive demo](interactive_demo) and include the [TVM](https://github.com/apache/tvm) tuned model. It achieves **8FPS** on Jetson Nano GPU now!

**[NEW!]** Add support to the [MUNIT](https://github.com/NVlabs/MUNIT), a multimodal unsupervised image-to-image translation approach! Please follow the [test commands](#munit) to test the pre-trained models and the [tutorial](docs/tutorials/fast_gan_compression.md) to train your own models!

![teaser](imgs/teaser.png)
*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](https://lmxyy.me/), [Ji Lin](http://linji.me/), [Yaoyao Ding](https://yaoyaoding.com/), [Zhijian Liu](http://zhijianliu.com/), [Jun-Yan Zhu](https://www.cs.cmu.edu/~junyanz/), and [Song Han](https://songhan.mit.edu/)

MIT, Adobe Research, SJTU

In CVPR 2020.

## Demos

## Overview

![overview](imgs/overview.png)*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

![performance](imgs/performance.jpeg)

*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](https://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/cycle_gan.ipynb) and [pix2pix](https://colab.research.google.com/github/mit-han-lab/gan-compression/blob/master/pix2pix.ipynb).

## Prerequisites

* Linux
* Python 3
* CPU or NVIDIA GPU + CUDA CuDNN

## Getting Started

### Installation

- Clone this repo:

```shell
git clone [email protected]:mit-han-lab/gan-compression.git
cd gan-compression
```

- Install [PyTorch](https://pytorch.org) 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).

```shell
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,

```shell
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.

```shell
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.

```shell
bash scripts/cycle_gan/horse2zebra/test_full.sh
```

* Test the compressed model.

```shell
bash scripts/cycle_gan/horse2zebra/test_compressed.sh
```
* Measure the latency of the two models.

```shell
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:

```bash
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).

```shell
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,

```shell
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.

```shell
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.

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

* Test the compressed model.

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

* Measure the latency of the two models.

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

```shell
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](#cityscapes) 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,

```shell
bash datasets/download_real_stat.sh cityscapes A
```

#### Apply a Pre-trained Model

* Download the pre-trained models.

```shell
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.

```shell
bash scripts/gaugan/cityscapes/test_full.sh
```

* Test the compressed model.

```shell
bash scripts/gaugan/cityscapes/test_compressed.sh
```
* Measure the latency of the two models.

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

```shell
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).

```shell
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,

```shell
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.

```shell
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.

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

* Test the compressed model.

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

* Measure the latency of the two models.

```shell
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](datasets/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:

```shell
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](https://github.com/NVlabs/SPADE). Specifically, you need to download `train2017.zip`, `val2017.zip`, `stuffthingmaps_trainval2017.zip`, and `annotations_trainval2017.zip` from [nightrome/cocostuff](https://github.com/nightrome/cocostuff). The images, labels, and instance maps should be arranged in the same directory structure as in [datasets/coco_stuff](https://github.com/NVlabs/SPADE/tree/master/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](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](https://github.com/kazuto1011/deeplab-pytorch/releases/download/v1.0/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

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

Fast GAN Compression
26.0M
35.5G
25.06
35.05

MUNIT
edges→shoes
Original
15.0M
77.3G
30.13
--

Fast GAN Compression
1.10M
2.63G
30.53
--

### Training

Please refer to the tutorial of [Fast GAN Compression](docs/tutorials/fast_gan_compression.md) and [GAN Compression](docs/tutorials/gan_compression.md) 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](./get_real_stat.py) to extract statistical information. For example, for the edges2shoes dataset, you could run the following command:

```shell
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](docs/overview.md)

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](https://arxiv.org/pdf/2003.08936).
```bibtex
@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](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix), [SPADE](https://github.com/NVlabs/SPADE), and [MUNIT](https://github.com/NVlabs/MUNIT).

We also thank [pytorch-fid](https://github.com/mseitzer/pytorch-fid) for FID computation, [drn](https://github.com/fyu/drn) for cityscapes mIoU computation and [deeplabv2](https://github.com/kazuto1011/deeplab-pytorch) for Coco-Stuff mIoU computation.