Applying Waseerstein GAN to SRGAN, a GAN based super resolution algorithm.
This repo was forked from @zsdonghao 's tensorlayer/srgan repo, based on this original repo, I changed some code to apply wasserstein loss, making the training procedure more stable, thanks @zsdonghao again, for his great reimplementation.
TensorFlow Implementation of "Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network"
When the SRGAN was first proposed in 2016, we haven't had Wasserstein GAN(2017) yet, WGAN using wasserstein distance to measure the disturibution difference between different data set. As for the original GAN training, we don't know when to stop training the discriminator or the generator, to get a nice result. But when using the wasserstein loss, as the loss decreasing, the result will be better. So we are going to use the WGAN and we are not going to explain the math detail of WGAN here, but to give the following steps to apply WGAN.
model.py
, line 218-219)main.py
, line 105-108)main.py
, line 136)main.py
, line 132-133)These above steps was given by an excellent article[4], the arthor explained the WGAN in a very straightforward way, it was written in Chinese.
config.py
(like number of epochs) are seleted basic on that dataset, if you change a larger dataset you can reduce the number of epochs. train_hr_imgs = tl.files.load_flickr25k_dataset(tag=None)
in main.py
. config.TRAIN.hr_img_path
in config.py
.We run this script under TensorFlow 1.4 and the TensorLayer 1.8.0+.
pip install tensorlayer==1.8.0
conda install tensorflow-gpu==1.3.0
pip install tensorflow-gpu==1.4.0
pip install easydict
config.py
. config.TRAIN.img_path = "your_image_folder/"
I added the tensorboard callbacks to monitor the training procedure, please change the logdir to your folder.
config.VALID.logdir = 'your_tensorboard_folder'
python main.py
SRGAN_g
(model.py line 53) is using 1×1 kernel, but I changed this kernel to 9×9, so if you use this pretrained weights, you may get the weights unequal error.
Two advice:
1)Train the whole network from scratch, you'll get the 9×9 version weights, for further training or evaluating images.
2)You can just change the SRGAN_g
's final conv kernel (model.py
line 53) to (1, 1) instead of (9, 9), and change the model.py
line 35 conv kernel from (9, 9) to (3, 3), so that you can use the pretrained weights.python main.py --mode=evaluate
Compare with the original version, I did the following changes:
model.py
(line 100), changing the kernel size from (1, 1) to (9, 9), as the paper proposed.