yuantian013 / e2gan

[ECCV 2020]"Off-Policy Reinforcement Learning for Efficient and Effective GAN Architecture Search" By Yuan Tian, Qin Wang, Zhiwu Huang, Wen Li, Dengxin Dai, Minghao Yang, Jun Wang, Olga Fink

E2GAN

Code for Off-Policy Reinforcement Learning for Efficient and Effective GAN Architecture Search, ECCV 2020.

Introduction

We formulate the GAN architecture search problem as a Markov decision process (MDP) inspired by the success of human-designed Progressive GAN. This new formulation enables us to discover competitive GAN architectures on a single 2080TI in 7 hours using off-policy RL.

Dependencies

conda create --name e2ganrl python=3.6
conda activate e2ganrl

conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.0 -c pytorch

python3 -m pip install imageio
python3 -m pip install scipy
python3 -m pip install six
python3 -m pip install numpy==1.18.1
python3 -m pip install python-dateutil==2.7.3
python3 -m pip install tensorboardX==1.6
# For the reward calculation, external tf code
python3 -m pip install tensorflow-gpu==1.13.1
python3 -m pip install tqdm==4.29.1

Code was tested on a RTX2080TI with 11GB RAM.

Prepare fid statistic file

Download the pre-calculated statistics from AutoGAN (Link) to ./search/fid_statand ./eval/fid_stat .

Run E2GAN search on CIFAR-10

cd search
bash exps/search.sh

You will find the architectures in the log file ./search/search.log after running the above script.

Train from scratch the discovered architecture

To train from scratch and get the performance of your discovered architecture, run the following command (you should replace the architecture vector following "--arch" in the script with best-performing candidate architectures in the exploitation stage in search.log):

cd eval
# Train the discovered GAN on CIFAR-10
bash exps/train_derived.sh
# Train the discovered GAN on STL
bash exps/train_derived_stl.sh

Test the architecture reported in the paper

Run the following script:

cd eval
# Testing the pretrained CIFAR-10 Model
bash exps/test.sh
# Testing the pretrained STL Model
bash exps/test_stl.sh

Pre-trained models (both CIFAR and STL) are provided (Google Drive). Please put them in eval/checkpoints/ .

Citation

Please cite our work if you find it useful.

@InProceedings{Tian_2020_ECCV,
author = {Yuan Tian, Qin Wang, Zhiwu Huang, Wen Li, Dengxin Dai, Minghao Yang, Jun Wang, Olga Fink},
title = {Off-Policy Reinforcement Learning for Efficient and Effective GAN Architecture Search},
booktitle = {The European Conference on Computer Vision (ECCV)},
month = {September},
year = {2020}
}

Supplementary material

PDF

Acknowledgement

1. Inception Score code from OpenAI's Improved GAN (official).
2. FID code and CIFAR-10 statistics file from https://github.com/bioinf-jku/TTUR (official).
3. SAC code from https://github.com/pranz24/pytorch-soft-actor-critic.
4. GAN training/eval code is heavily borrowed from AutoGAN https://github.com/TAMU-VITA/AutoGAN

For questions regarding the code, please open an issue or contact Yuan and Qin via email {yutian, qwang} AT ethz.ch

e2gan's Issues

Training on 512x512

I tried to apply your code to a 512x512x3 set of images, in a real world dataset.
I made several modifications to the basic code for data-loading and used the cal_fid_stat.py script from AutoGAN to generate stats for my test-set.

I modified the search.sh file to the following:

CUDA_VISIBLE_DEVICES=3 python -u search.py \
-gen_bs 16 \
-dis_bs 8 \
--dataset stl10 \
--bottom_width 4 \
--img_size 512 \
--gen_model shared_gan \
--dis_model shared_gan \
--controller controller \
--latent_dim 512 \
--gf_dim 512 \
--df_dim 256 \
--g_spectral_norm False \
--d_spectral_norm True \
--g_lr 0.0002 \
--d_lr 0.0002 \
--beta1 0.0 \
--beta2 0.9 \
--init_type xavier_uniform \
--n_critic 5 \
--val_freq 20 \
--ctrl_sample_batch 1 \
--shared_epoch 15 \
--grow_step1 15 \
--grow_step2 35 \
--max_search_iter 65 \
--ctrl_step 30 \
--random_seed 12345 \
--exp_name e2gan_search --data_path /home/user/data-E2GAN | tee search.log

I run it and I got the the following error:

search progress:   0%|                                  | 0/100 [00:35<?, ?it/s]
Traceback (most recent call last):
  File "search.py", line 227, in <module>
    main()
  File "search.py", line 155, in main
    action = Agent.select_action([layer, last_R,0.01*last_fid] + last_state,Best)
  File "/home/user/E2GAN/search/sac.py", line 60, in select_action
    action1,action2,action3,action4, action5,action6,_,_,_,_,_, _ ,_,_, _, _ ,_,_,= self.policy.sample(state)
  File "/home/user/E2GAN/search/sac_model.py", line 117, in sample
    mean_1, log_std_1,mean_2, log_std_2,mean_3, log_std_3,mean_4, log_std_4,mean_5, log_std_5,mean_6, log_std_6= self.forward(state)
  File "/home/user/E2GAN/search/sac_model.py", line 84, in forward
    x = F.relu(self.linear1(state.cuda()))
  File "/home/user/miniconda3/envs/ganspace/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/miniconda3/envs/ganspace/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 87, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/user/miniconda3/envs/ganspace/lib/python3.7/site-packages/torch/nn/functional.py", line 1370, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: size mismatch, m1: [1 x 515], m2: [131 x 128] at /opt/conda/conda-bld/pytorch_1573049306803/work/aten/src/THC/generic/THCTensorMathBlas.cu:290

I guess the error has to do with downscaling/upscaling convolutions, but I am not sure.
I was curious if you had tried a 512px model in the past or if there is a straightforward problem you can observe in the configuration of my script file.

Thanks in advance!