Skip to content

Latest commit

 

History

History
68 lines (46 loc) · 2.26 KB

README.md

File metadata and controls

68 lines (46 loc) · 2.26 KB
Raw

Image-to-image translation with flow-based generative model

Requirements

  • Tensorflow (tested with v1.8.0)
  • Horovod (tested with v0.13.8) and (Open)MPI

Run

pip install -r requirements.txt

To setup (Open)MPI, check instructions on Horovod github page.

Dataset

cd pix2pix-flow
  • edges2shoes: Download the dataset here.
  • mnist: Nothing need to be done.

Train

Train with 4 GPUs:

mpiexec -n 4 python train.py --problem PROBLEM --image_size 32 --n_level 3 --depth 32 --flow_permutation 2 --flow_coupling 1 --seed 0 --learntop --lr 0.001 --n_bits_x 8 --joint_train --logdir LOGDIR

Replace PROBLEM with [mnist | edges2shoes], and LOGDIR with the repo you want to save the trained models.

Inference

python train.py --problem PROBLEM --image_size 32 --n_level 3 --depth 32 --flow_permutation 2 --flow_coupling 1 --seed 0 --learntop --lr 0.001 --n_bits_x 8 --joint_train --logdir LOGDIR --inference

Replace PROBLEM with [mnist | edges2shoes], and and LOGDIR with the repo you want to save the inference results.

After running the command, you will get:

  • z_A.npy: Latent code of images in domain A from test set.
  • z_B.npy: Latent code of images in domain B from test set.
  • A2B.png: Images in domain B translated from domain A.
  • B2A.png: Images in domain A tanslated from domain B.

in LOGDIR.

DEBUG: train with 1 GPU

Run wtih small depth to test

CUDA_VISIBLE_DEVICES=0 python train.py --depth 1 --epochs_full_sample 1 --epochs_full_valid 1 --joint_train --problem PROBLEM --logdir LOGDIR

Work in Progress

Train with 4 GPUs with B domain pixel l2 loss:

mpiexec -n 4 python train.py --problem PROBLEM --image_size 32 --n_level 3 --depth 32 --flow_permutation 2 --flow_coupling 1 --seed 0 --learntop --lr 0.001 --n_bits_x 8 --joint_train --logdir LOGDIR --B_loss --B_loss_fn l2

Disclaimer

The code is hugely borrowed from OpenAI's Glow.