liannice/gan-slimming

[ECCV 2020] "All-in-One GAN Compression by Unified Optimization" by Haotao Wang, Shupeng Gui, Haichuan Yang, Ji Liu, and Zhangyang Wang

GAN-Slimming

GAN Slimming: All-in-One GAN Compression by A Unified Optimization Framework

Haotao Wang, Shupeng Gui, Haichuan Yang, Ji Liu, Zhangyang Wang

In ECCV 2020 (Spotlight)

Overview

An all-in-one GAN compression method integrating model distillation, channel pruning and quantization under GAN minimax optimization framework.

Visualization Results

Image-to-image translation by (compressed) CycleGAN:

Training

1. Download dataset:

./download_dataset <dataset_name>

This will download the dataset to folder datasets/<dataset_name> (e.g., datasets/summer2winter_yosemite).

2. Train origianl dense CycleGAN and generate style stransfer results on training set:

Use the offcial CycleGAN codes to train origianl dense CycleGAN and generate style stransfer results on training set. Put the style transfer results to folder train_set_result/<dataset_name>. For example, train_set_result/summer2winter_yosemite/B/2009-12-06 06:58:39_fake.png is the fake winter image transfered from the real summer image datasets/summer2winter_yosemite/A/2009-12-06 06:58:39.png using the orignal dense CycleGAN.

3. Compress

GS-32:

python gs.py --rho 0.01 --dataset <dataset_name> --task <task_name>

GS-8:

python gs.py --rho 0.01 --quant --dataset <dataset_name> --task <task_name>

The training results (checkpoints, loss curves, etc.) will be saved in results/<dataset_name>/<task_name>. Valid <dataset_name>s are: horse2zebra, summer2winter_yosemite. Valid <task_name>s are: A2B, B2A. (For example, horse2zebra/A2B means transferring horse to zebra and horse2zebra/B2A means transferring zebra to horse.)

4. Extract compact subnetwork obtained by GS

GAN slimming has pruned some channels in the network by setting the channel-wise mask to zero. Now we need to extract the actual compressed subnetowrk.

python extract_subnet.py --dataset <dataset_name> --task <task_name> --model_str <model_str> 

The extracted subnetworks will be saved in subnet_structures/<dataset_name>/<task_name>

5. Finetune subnetwork

python finetune.py --dataset <dataset_name> --task <task_name> --base_model_str <base_model_str>

Finetune results will be saved in finetune_results/<dataset_name>/<task_name>

Pretrianed Models

Pretrained models are available through Google Drive.

Citation

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

@inproceedings{wang2020ganslimming,
  title={GAN Slimming: All-in-One GAN Compression by A Unified Optimization Framework},
  author={Wang, Haotao and Gui, Shupeng and Yang, Haichuan and Liu, Ji and Wang, Zhangyang},
  booktitle={European Conference on Computer Vision},
  year={2020}
}

Our Related Work

Please also check our concurrent work on combining neural architecture search (NAS) and model distillation for GAN compression:

Yonggan Fu, Wuyang Chen, Haotao Wang, Haoran Li, Yingyan Lin, and Zhangyang Wang. "AutoGAN-Distiller: Searching to Compress Generative Adversarial Networks." ICML, 2020. [pdf] [code]