Skip to content
master
Go to file
Code

Files

Permalink
Failed to load latest commit information.
Type
Name
Latest commit message
Commit time
img
Jun 15, 2020
Jun 14, 2020

README.md

AutoGAN-Distiller: Searching to Compress Generative Adversarial Networks

Yonggan Fu, Wuyang Chen, Haotao Wang, Haoran Li, Yingyan Lin, Zhangyang Wang

Accepted at ICML 2020 [Paper Link].

Overview

We propose AutoGAN-Distiller (AGD) Framework, among the first AutoML frameworks dedicated to GAN compression, and is also among a few earliest works that explore AutoML for GANs.

Method

  • AGD is established on a specifically designed search space of efficient generator building blocks, leveraging knowledge from state-of-the-art GANs for different tasks.
  • It performs differentiable neural architecture search under the target compression ratio (computational resource constraint), which preserves the original GAN generation quality via the guidance of knowledge distillation.
  • We demonstrate AGD on two representative mobile-based GAN applications: unpaired image translation (using a CycleGAN), and super resolution (using an encoder-decoder GAN).

Visualization Results

Unpaired image translation:

unpair-image-translation

Super Resolution:

super-resolution

Datasets

Unpaired Image Translation

horse2zebra, zebra2horse, summer2winter, winter2summer: Unpaired-dataset

Super Resolution

Training (DIV2K+Flickr2K): SR-training-dataset

Evaluation (Set5, Set14, BSD100, Urban100): SR-eval-dataset

Usage

Overview

AGD_ST and AGD_SR are the source codes for unpaired image translation task and super resolution task respectively. The codes for pretrain, search, train from scratch and eval are in the AGD_ST/search and AGD_SR/search directory.

We use AGD_ST/search as an example. All the configurations during pretrain, search, train from scratch, eval are in config_search.py, config_train.py and config_eval.py respectively. Please specify the target dataset C.dataset and change the dataset path C.dataset_path in the three config files to the real paths on your PC.

Prerequisites

See env.yml for the complete conda environment. Create a new conda environment:

conda env create -f env.yml
conda activate pytorch

In partiqular, if the thop package encounters some version conflicts, please specify the thop version:

pip install thop==0.0.31.post1912272122

Step 1: Pretrain the Supernet

  • Switch to the search directory:
cd AGD_ST/search
  • Set C.pretrain = True in config_search.py.

  • Start to pretrain:

python train_search.py

The checkpoints during pretraining are saved at ./ckpt/pretrain.

Step 2: Search

  • Set C.pretrain = 'ckpt/pretrain' in config_search.py.

  • Start to search:

python train_search.py

Step 3: Train the derived network from scratch

  • Set C.load_path = 'ckpt/search' in config_train.py.

  • Start to train from scratch:

python train.py

Step 4: Eval

  • Set C.load_path = 'ckpt/search' and C.ckpt = 'ckpt/finetune/weights.pt' in config_eval.py.
  • Start to evaluate on the testing dataset:
python eval.py

The result images are saved at ./output/eval/.

Two differences in Super Resolution tasks

1st Difference

Please download the checkpoint of original ESRGAN (teacher model) from pretrained ESRGAN and move it to the directory AGD_SR/search/ESRGAN/.

2nd Difference

The step 3 is splitted into two steps, i.e., first pretrain the derived architecture with only content loss and then finetune with perceptual loss:

  • Pretrain: Set C.pretrain = True in config_train.py.

  • Finetune: Set C.pretrain = 'ckpt/finetune_pretrain/weights.pt' in config_train.py.

Pretrained Models

Pretrained models are provided at pretrained AGD.

To evaluate the pretrained models, please copy the network architecture definition and pretrained weights to the corresponding directories:

cp arch.pt ckpt/search/
cp weights.pt ckpt/finetune/

then do the evaluation following step 4.

Our Related Work

Please also check our concurrent work on a unified optimization framework combining model distillation, channel pruning and quantization for GAN compression:

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

About

[ICML2020] "AutoGAN-Distiller: Searching to Compress Generative Adversarial Networks" by Yonggan Fu, Wuyang Chen, Haotao Wang, Haoran Li, Yingyan Lin, Zhangyang Wang

Topics

Resources

License

Releases

No releases published

Packages

No packages published

Languages

You can’t perform that action at this time.