• Stars
    star
    1,104
  • Rank 42,038 (Top 0.9 %)
  • Language
    Python
  • License
    Other
  • Created over 4 years ago
  • Updated 6 months ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

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

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!

teaser 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

overviewGAN 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

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 [email protected]: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
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 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.

More Repositories

1

streaming-llm

[ICLR 2024] Efficient Streaming Language Models with Attention Sinks
Python
6,530
star
2

bevfusion

[ICRA'23] BEVFusion: Multi-Task Multi-Sensor Fusion with Unified Bird's-Eye View Representation
Python
2,286
star
3

temporal-shift-module

[ICCV 2019] TSM: Temporal Shift Module for Efficient Video Understanding
Python
2,060
star
4

once-for-all

[ICLR 2020] Once for All: Train One Network and Specialize it for Efficient Deployment
Python
1,866
star
5

llm-awq

AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
Python
1,687
star
6

proxylessnas

[ICLR 2019] ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware
C++
1,420
star
7

torchquantum

A PyTorch-based framework for Quantum Classical Simulation, Quantum Machine Learning, Quantum Neural Networks, Parameterized Quantum Circuits with support for easy deployments on real quantum computers.
Jupyter Notebook
1,304
star
8

data-efficient-gans

[NeurIPS 2020] Differentiable Augmentation for Data-Efficient GAN Training
Python
1,277
star
9

efficientvit

EfficientViT is a new family of vision models for efficient high-resolution vision.
Python
1,218
star
10

torchsparse

[MICRO'23, MLSys'22] TorchSparse: Efficient Training and Inference Framework for Sparse Convolution on GPUs.
Cuda
1,181
star
11

smoothquant

[ICML 2023] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
Python
1,175
star
12

anycost-gan

[CVPR 2021] Anycost GANs for Interactive Image Synthesis and Editing
Python
778
star
13

tinyml

Python
755
star
14

TinyChatEngine

TinyChatEngine: On-Device LLM Inference Library
C++
730
star
15

tinyengine

[NeurIPS 2020] MCUNet: Tiny Deep Learning on IoT Devices; [NeurIPS 2021] MCUNetV2: Memory-Efficient Patch-based Inference for Tiny Deep Learning; [NeurIPS 2022] MCUNetV3: On-Device Training Under 256KB Memory
C
717
star
16

fastcomposer

[IJCV] FastComposer: Tuning-Free Multi-Subject Image Generation with Localized Attention
Python
644
star
17

pvcnn

[NeurIPS 2019, Spotlight] Point-Voxel CNN for Efficient 3D Deep Learning
Python
639
star
18

lite-transformer

[ICLR 2020] Lite Transformer with Long-Short Range Attention
Python
589
star
19

spvnas

[ECCV 2020] Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution
Python
577
star
20

distrifuser

[CVPR 2024 Highlight] DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models
Python
538
star
21

mcunet

[NeurIPS 2020] MCUNet: Tiny Deep Learning on IoT Devices; [NeurIPS 2021] MCUNetV2: Memory-Efficient Patch-based Inference for Tiny Deep Learning
Python
460
star
22

tiny-training

On-Device Training Under 256KB Memory [NeurIPS'22]
Python
432
star
23

amc

[ECCV 2018] AMC: AutoML for Model Compression and Acceleration on Mobile Devices
Python
428
star
24

dlg

[NeurIPS 2019] Deep Leakage From Gradients
Python
400
star
25

haq

[CVPR 2019, Oral] HAQ: Hardware-Aware Automated Quantization with Mixed Precision
Python
368
star
26

offsite-tuning

Offsite-Tuning: Transfer Learning without Full Model
Python
365
star
27

hardware-aware-transformers

[ACL'20] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
Python
321
star
28

litepose

[CVPR'22] Lite Pose: Efficient Architecture Design for 2D Human Pose Estimation
Python
304
star
29

inter-operator-scheduler

[MLSys 2021] IOS: Inter-Operator Scheduler for CNN Acceleration
C++
191
star
30

amc-models

[ECCV 2018] AMC: AutoML for Model Compression and Acceleration on Mobile Devices
Python
166
star
31

apq

[CVPR 2020] APQ: Joint Search for Network Architecture, Pruning and Quantization Policy
Python
156
star
32

parallel-computing-tutorial

C++
134
star
33

flatformer

[CVPR'23] FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer
Python
119
star
34

patch_conv

Patch convolution to avoid large GPU memory usage of Conv2D
Python
74
star
35

6s965-fall2022

Jupyter Notebook
64
star
36

sparsevit

[CVPR'23] SparseViT: Revisiting Activation Sparsity for Efficient High-Resolution Vision Transformer
Python
48
star
37

bnn-icestick

Binary Neural Network on IceStick FPGA.
Jupyter Notebook
47
star
38

e3d

Efficient 3D Deep Learning
46
star
39

neurips-micronet

[JMLR'20] NeurIPS 2019 MicroNet Challenge Efficient Language Modeling, Champion
Jupyter Notebook
40
star
40

spatten-llm

[HPCA'21] SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning
Scala
32
star
41

tinychat-tutorial

C++
28
star
42

pruning-sparsity-publications

14
star
43

iccad-tinyml-open

[ICCAD'22 TinyML Contest] Efficient Heart Stroke Detection on Low-cost Microcontrollers
C
14
star
44

calo-cluster

Jupyter Notebook
5
star
45

ml-blood-pressure

Python
5
star
46

gan-compression-dynamic

Python
3
star
47

data-efficient-gans-dynamic

Python
3
star