Awesome
COCO-GAN
Authors' implementation of "COCO-GAN: Generation by Parts via Conditional Coordinating" in TensorFlow.
[Project Page] [Paper] [Paper (Full Resolution)]
0. Pre-requirements
# Install the major required packages with
conda install tensorflow-gpu==1.13.1 scipy==1.1.0 pyyaml==5.3 pillow==6.0.0 tqdm
# Note:
# Recommended Tensorflow versions: 1.9.*~1.13.*
# Other Tensorflow versions may be applicable, but you may face some errors during the FID calculation.
You may directly disable FID calculation, or try to comment/uncomment different variants of TensorShape
setting codes I used for different TF versions at fid_utils/fid.py:86~89
.
You may need to write your own workaround if you face errors with other TF versions.
1. Data Download and Preprocessing
CelebA
You may directly run the following command to setup everythong:
python scripts/download_celeba.py
LSUN (Bedroom)
You may directly run the following command to setup everythong:
sh scripts/download_lsun_bedroom.sh
Matterport3D
- You must first contact the authors of Matterport3D dataset, sign an user license agreement, then obtain the downaload script from the authors.
- After you obtained the
download_mp.py
, download thematterport_skybox_images
category and place it under./data/matterport3d/
. Please refer to the official document:https://github.com/niessner/Matterport/blob/master/data_organization.md
. - Install cube2sphere library with
pip install cube2sphere
. - Execute the preprocessing script with
python scripts/preprocessing_mp3d.py
, which outputs the final images at./data/matterport3d_panorama/
.
Note: I don't have access to the original dataset now, so I'm not entirely sure about the detailed filenames and paths. You may need to make some small modifications in scripts/preprocessing_mp3d.py:14:16
.
CelebA-syn
A synthesis variant of CelebA that mess-ups the alignment of faces in CelebA dataset.
python ./scripts/generate_celeba_syn --data_dir="./data/CelebA/*.jpg" --syn_type="inward"
python ./scripts/generate_celeba_syn --data_dir="./data/CelebA/*.jpg" --syn_type="outward"
CelebA-HQ
I forgot how I downloaded it. I believe there was a script somewhere though...
Please manually download it from the official site: https://github.com/tkarras/progressive_growing_of_gans
2. Convert TF-Records
To speed up training, we process the images into tfrecords:
# Run the one(s) you need
python ./scripts/compute_tfrecord.py --dataset celeba --resolution 64
python ./scripts/compute_tfrecord.py --dataset celeba --resolution 128
python ./scripts/compute_tfrecord.py --dataset lsun --resolution 64
python ./scripts/compute_tfrecord.py --dataset lsun --resolution 256
python ./scripts/compute_tfrecord.py --dataset matterport3d
python ./scripts/compute_tfrecord.py --dataset celeba-hq
If your data is in a directory other than ./data/
, another way to do the conversion is:
python ./scripts/compute_tfrecord.py --dataset celeba --resolution 64 --img_paths="./path/to/images/*.jpg"
3. Pre-calculate FID statistics
Calculation of FID requires some statistic from real data, you may disable FID calcuation in the config files.
python ./fid_utils/precalc_fid_stats.py --dataset lsun --data_path "./data/lsun/*"
This will generate statistic files under ./stats/
.
4. Training
We provide the configurations for standard COCO-GAN training on CelebA, LSUN and Matterport3D:
python main.py --config="./configs/CelebA_128x128_N2M2S64.yaml"
python main.py --config="./configs/LSUN_256x256_N2M2S128.yaml"
python main.py --config="./configs/MP3D_258x768_N2M2S128.yaml"
Note: N,M
is the number of micro patches in a macro patch, S
is the macro patch size.
CelebA-HQ is not well-tested (actually, we only tune the hyperparameters for CelebA):
python main.py --config="./configs/CelebAHQ_1024x1024_N2M2S512.yaml"
Also the extrapolation experiment (this requires a well-pretrained LSUN_256x256_N2M2S128
checkpoint):
python main.py --config="./configs/LSUN_256x256_N2M2S128_Extrapolation.yaml"
Note: This experiment is a little unstable, I usually run 3 to 5 epochs and pick the best model by personal preference.
With the Patch-Guided Image Generation, we need to train an additional critic Q
by setting Q_update_period=1
and code_loss_w=100
. Here's an exmple config:
python main.py --config="./configs/CelebA_128x128_N2M2S64_PatchGuidedGeneration.yaml"
5. Running Testing Only
(Not well-tested) (You can skip step 1~5)
Suppose you have successfully adopted either the default path (specified in the yaml files) to save checkpoints or use force_load_from_dir
to find checkpoints, you can run the following commands:
python main.py --config="./configs/CelebA_128x128_N2M2S64.yaml" --test
The testing codes also provides two arguments: n_samples
and test_output_dir
.
Pretrained checkpoints
We provide following pretrained checkpoints:
- CelebA_64x64_N2M2S32
- CelebA_128x128_N2M2S64
- CelebA_128x128_N2M2S64_PatchGuidedGeneration
- LSUN_64x64_N2M2S32
- LSUN_256x256_N2M2S128
Please download them from: https://drive.google.com/drive/folders/1Mr5BknOrTebQgxdARxJVV95pZA-NkVXt?usp=sharing
Then, you can use force_load_from_dir: "path/to/the/pretrained/directory"
argument in each of the config file to load the parameters (this does not override the hyperparameters).
Note:
- These models are trained with Tensorflow 1.13.0, you may potentially face some loading errors if you use other Tensorflow versions.
- The structure of the checkpoint directory is:
./logs/
+ <some exp A> (name that matches the `exp_name` in the yaml file, e.g., `CelebA_64x64_N2M2S32`)
| +---ckpt/
| | +--- (place the downloaded checkpoint files here)
| +---images/
| +---(event logs)
|
+ <some exp B>
|
.
.
Some Tuning Tricks While Training on Custom Datasets
(In cases you received NaN exceptions after running with the default hyperparameters)
As you have seen, the adversarial loss values are quite magnificent in COCO-GAN training (from 1e4 to 1e8, depending on the complexity of the images in the dataset). The training is sometimes unstable. And gradient penalty plays an important role in stabilizing the training. And yes, you may need to tune the gradient penalty strength a bit if there are evidences of training instability, for instance, NaNs LOL.
In order to determine the appropriate gradient penalty strength, please open the tensorboard and look into your previous training logs. Go to histogram and open the gp_slopes
pane. Long story short, we would like the constrain the mean of the distribution to near 1, which is also known as the 1-Lipchitz constraint. You may start to increase the gp_lambda
in the config to pull the distribution toward zero a bit, just a bit, otherwise, your training will saturate when the gradients are way too small. Note that, in most cases, you only need to make sure the mean of gp_lambda
is around 1~5 after 10K iterations, as the early stage is expectedly quite unstable.
Some Additional Notes
- Some computers take forever to calculate FID, it is usually an unknown problem in the CPU, please switch a computer or disable FID calculation!
- Though our implementation theretically supports (a) patch size H!=W, and (b) micro patch N!=M cases, but these functions are never tested. Please let me know (or create a pull request with a fix if you are willing) if you face any problem while using this feature!
- In general, with smaller micro patches, as the computation graph becomes too complex for Tensorflow, it will take lots of time an GPU memory to build the graph. The performance may be improved with Pyotrch and using
torch.no_grad()
smartly. - This implementation is different to our private codebase, please kindly let me know if you found anything bizzard.
- The coordinate generation part looks complicated since I made it generic to different coordinate designs.
Special Thanks
The following open-source repositories largely facilitate our research!
- Spectral Normalization: https://github.com/minhnhat93/tf-SNDCGAN
- cGANs with Projection Discriminator (we replicate its architecture): https://github.com/pfnet-research/sngan_projection
- FID: https://github.com/bioinf-jku/TTUR
- Gradient Penalty: https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb
- Some GAN-Utilities: https://github.com/carpedm20/DCGAN-tensorflow
- Cube to Sphere transform: https://github.com/Xyene/cube2sphere
Citation
@inproceedings{lin2019cocogan,
author = {Lin, Chieh Hubert and
Chang, Chia{-}Che and
Chen, Yu{-}Sheng and
Juan, Da{-}Cheng and
Wei, Wei and
Chen, Hwann{-}Tzong},
title = {{COCO-GAN:} Generation by Parts via Conditional Coordinating},
booktitle = {IEEE International Conference on Computer Vision (ICCV)},
year = {2019},
}