yueying-teng / pro_gan_tf2

Progressively grow faces using CELEBA HQ data

TensorFlow 2 Implementation of Progressive Growing of GANs

📄Paper

How to use

1. Build and run the container

bash bin/docker_build.sh

bash bin/docker_run.sh

2. Prepare the data

• For this project, the training set (both male and female) of CELEBA HQ data is used, and it's downloaded from Kaggle.
• Put the downloaded data at the root dir of the project

3. Train the models

bash bin/log_into_container.sh
cd notebooks

# run the training script; logs will be exported to anohter file
python3 train_pro_gan.py --data_dir "/work/CELEBAHQ/train" --save_dir "./test_train" --epochs 40 > pro_gan_training.log 2>&1

4. Start tensorboard

# open a new terminal window
bash bin/log_into_container.sh

# NOTE that the log dir below should be replaced with the save_dir/logs you specified
bash bin/start_tensorboard.sh /work/notebooks/test_train/logs

Project structure

- pro_gan_tf
    - bin
    - CELEBAHQ  <-- replace this with your data
    - config
    - demo
    – docker
    - notebooks   <-- train_pro_gan.py is here, logs are also saved here while training
        - train_pro_gan.py
        - test_train
            - logs/20220105-142028
            - models/20220105-142028
    - pro_gan
    - README.md
    - requirements.txt
    - test
    - training_results

To create demo videos

# open a new terminal window
bash bin/log_into_container.sh
cd demo

# create a video using the images saved during training
python3 create_video_from_training_feedback.py --log_dir "/work/notebooks/test_train/logs/20220105-142028" --depth 7

# create a video using the images generated from the interpolated latent space vectors
python3 latent_space_interpolation.py --model_dir "/work/notebooks/test_train/models/20220105-142028" --video_length 10 --depth 7 --num_videos 4

# use the trained generator model to create a sheet of generated images
python3 create_facesheet.py --model_dir "/work/notebooks/test_train/models/20220105-142028" --num_samples 16 --depth 7 --num_sheets 4

Some training results

• One Tesla P40 is used for training the models up till resolution 128 x 128 before running out of memory.
• This training process took about 25 days and smaller batch sizes are used for training models with image resolution 64 and above.

Latent space interpolation

Facesheet 128 x 128

   

Generator model w/ resolution 128 x 128

Click here to download the model.

To generate facesheets and latent space interpolation videos, put the downloaded model at notebooks/test_train/models/.

References:

• akanimax/pro_gan_pytorch
• tkarras/progressive_growing_of_gans
• tensorflow_gan/python/losses/losses_impl
• tensorflow_gan/examples/progressive_gan

TODOs:

• add support for
  • conditional GAN
  • continue training from any depth
  • EMA update of generator weights
• add metric FID