Gan produces desaturated, faded images - tensorflow

Question

I have a tensorflow implementation of a GAN that I'm trying to use to sharpen the image output from a variational autoencoder. The problem is after a few epochs (50000 generator weight updates, 200000 critic weight updates, single image batches) the network is still producing images that are almost greyscale when being trained on color images, and will typically produce very faded images with little contrast (although will produce clear images at times)
The generator gets a blurry reconstruction image as an input and tries to make a convincing fake of the real image it's based on. The discriminator gets a blurry reconstruction image and either the real image it's based on or the generators fake and ouputs a confidence value for if it is the real image or not.
This might be a case of not having trained the network enough but on every other GAN I've trained to slightly alter an image it normally learns to reproduce the input image within the first 300-400 weight updates.
One theory I have is that the gradient clipping is set quite low to facilitate stability (0.0001) and this might be causing a gradient saturation issue, but I honestly don't understand enough about the concepts involved to even know where to begin bugfixing this. This is the first time I've used wasserstein loss as well so let me know if my implementation is completely wrong.
Some additional information from my attempts to debug this -
I've tested the image inputs by adding mean absolute error loss to the generator and it is able to reproduce it's inputs and approximate the real images just fine
The loss is wasserstein and generally hangs around these values:
disc_total[0.000117] disc_real[-0.000080] disc_gen[0.000197] gen_total[-0.000062]
The input image folder is based on UTK face and contains images that look like this
The network output typically looks like this
Code:
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.utils import Sequence
from tensorflow.keras import backend as K
from tensorflow.keras.constraints import Constraint
import matplotlib.pyplot as plt
import vae
import numpy as np
import random
import cv2
import os
import time
#--------------------Build Model--------------------
#Save filepath
checkpoint_dir = 'Models/GAN/WGAN-U'
#Saved models filepath
filepath = "Models/BVAE/BVAE"
#Data filepath
datapath = "Datasets/Reconstructions/UTK"
val_datapath = ""
batch_size = 1
image_save_freq = 100
model_save_freq = 1000
grad_clipping = 0.0001
vae.vae.load_weights(filepath)
for layer in vae.vae.layers:
layer.trainable = False
#Wrapper for utk face data generator
class utk_data_gen(Sequence):
def __init__(self, validation):
self.validation = validation
self.dp = val_datapath if self.validation else datapath
self.data_len = len(os.listdir(self.dp))//batch_size
def __getitem__(self, index):
X1 = np.empty(shape=(batch_size, vae.height, vae.width, vae.channels))
X2 = np.empty(shape=(batch_size, vae.height, vae.width, vae.channels))
C = np.empty(shape=(batch_size, 7))
data_dir = os.listdir(self.dp)
i=index*batch_size
for j in range(0, batch_size):
data = data_dir[i+j]
img = cv2.imread(os.path.join(self.dp,data))
img = (img / 127.5)-1.0
vae_img = img[:vae.height]
img = img[vae.height:]
labels = data.split('_')
age = float(labels[0]) / 116
gender = float(labels[1])
white = 0.0
black = 0.0
asian = 0.0
indian = 0.0
other = 0.0
if labels[2] == '0':
white = 1.0
elif labels[2] == '1':
black = 1.0
elif labels[2] == '2':
asian = 1.0
elif labels[2] == '3':
indian = 1.0
elif labels[2] == '4':
other = 1.0
img_label = (age, gender, white, black, asian, indian, other)
img_label = np.array(img_label)
X1[j]=img
X2[j]=vae_img
C[j]=img_label
Y=X1
return [X1,X2,C],Y
def __len__(self):
return self.data_len
#Wasserstein discriminator
def wasserstein_loss(y_true, y_pred):
return K.mean(y_true * y_pred)
class ClipConstraint(Constraint):
def __init__(self, clip_value):
self.clip_value = clip_value
def __call__(self, weights):
return K.clip(weights, -self.clip_value, self.clip_value)
def get_config(self):
return {'clip_value': self.clip_value}
disc_optimizer = tf.keras.optimizers.RMSprop(lr=0.00005)
def build_discriminator():
init = tf.random_normal_initializer(0.0, 0.02)
kcon = ClipConstraint(grad_clipping)
img_input = layers.Input(shape=(vae.height, vae.width, vae.channels), name='input_image')
vae_input = layers.Input(shape=(vae.height, vae.width, vae.channels), name='vae_input_image')
concat_inputs = layers.Concatenate()([img_input, vae_input])
downsamples = 4
startFilters = 32
x = layers.Conv2D(startFilters, (4,4), strides=(1,1), padding='same', kernel_constraint=kcon, kernel_initializer=init)(concat_inputs)
for i in range(downsamples):
x = layers.Conv2D(startFilters * 2**i, (4,4), strides=(2,2), padding='same', kernel_constraint=kcon, kernel_initializer=init)(x)
x = layers.BatchNormalization()(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Conv2D(startFilters * 2**downsamples, (4,4), padding='same', kernel_constraint=kcon, kernel_initializer=init)(x)
x = layers.BatchNormalization()(x)
x = layers.LeakyReLU(alpha=0.2)(x)
x = layers.Flatten()(x)
output = layers.Dense(1, activation='linear')(x)
model = keras.Model([img_input, vae_input], output, name="discriminator")
model.compile(loss=wasserstein_loss, optimizer=disc_optimizer)
return model
#Generator
def UNet_downsample(layer_input, filters):
init = tf.random_normal_initializer(0.0, 0.02)
x = layers.Conv2D(filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_input)
x = layers.BatchNormalization()(x)
x = layers.LeakyReLU(alpha=0.2)(x)
return x
def UNet_upsample(layer_input, downsample_skip, filters, crop=False, use_dropout=False):
init = tf.random_normal_initializer(0.0, 0.02)
x = layers.Conv2DTranspose(filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_input)
if crop:
x = x[:,:-1,:-1]
x = layers.BatchNormalization()(x)
if use_dropout:
x = layers.Dropout(0.3)(x, training=True)
x = layers.Concatenate()([x, downsample_skip])
x = layers.LeakyReLU(alpha=0.2)(x)
return x
#U-Net
def build_generator():
init = tf.random_normal_initializer(0.0, 0.02)
input_img = keras.Input(shape=(vae.height, vae.width, vae.channels), name = "gen_input_image")
x = layers.Conv2D(3, (4,4), strides=1, padding="same", kernel_initializer=init, name="input_conv")(input_img)
x1 = UNet_downsample(x, 128)
x2 = UNet_downsample(x1, 256)
x3 = UNet_downsample(x2, 512)
x4 = UNet_downsample(x3, 512)
x5 = UNet_downsample(x4, 512)
x = layers.Conv2D(512, (4,4), strides=(2,2), padding='same', activation='relu', kernel_initializer=init)(x5)
x = UNet_upsample(x, x5, 512, True)
x = UNet_upsample(x, x4, 512, True, True)
x = UNet_upsample(x, x3, 512, True, True)
x = UNet_upsample(x, x2, 256)
x = UNet_upsample(x, x1, 128)
output = layers.Conv2DTranspose(3, (4,4), strides=(2,2), activation="tanh", padding="same", name="generator_outputs")(x)
return keras.Model(inputs = [input_img], outputs=output, name="generator")
#Build GAN
gan_optimizer = tf.keras.optimizers.RMSprop(lr=0.00005)
def build_gan(disc, gen):
for layer in disc.layers:
if not isinstance(layer, layers.BatchNormalization):
layer.trainable = False
input_img = keras.Input(shape=(vae.height,vae.width,vae.channels), name = "gan_input_img")
gen_output = gen(input_img)
disc_output = disc([gen_output, input_img])
model = keras.Model(input_img, [disc_output, gen_output])
model.compile(loss=[wasserstein_loss], optimizer=gan_optimizer)
return model
discriminator = build_discriminator()
generator = build_generator()
gan = build_gan(discriminator, generator)
generator_optimizer = tf.keras.optimizers.RMSprop(lr=0.00005)
discriminator_optimizer = tf.keras.optimizers.RMSprop(lr=0.00005)
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(gan_optimizer=gan_optimizer,
discriminator_optimizer=discriminator_optimizer,
discriminator=discriminator,
gan=gan)
cp_num = 0
def saveImg():
labels = np.array([[0.8, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]])
im = vae.decoder.predict([np.zeros(shape=(1,20)), labels])[0]
im = tf.expand_dims(im, axis=0)
im = generator.predict([im])[0]
im = ((im+1.0)*127.5).astype(np.uint8)
cv2.imwrite("Models/GAN/Progress/GAN_"+str(cp_num)+"_"+str(time.time())+ ".jpg", im)
mean_loss_index = 0
mean_loss_losses = [[0,]*30 for i in range(5)]
def get_mean_loss(gan_loss):
global mean_loss_index
global mean_loss_losses
for i, loss in enumerate(gan_loss):
mean_loss_losses[i][mean_loss_index] = loss
mean_loss_index += 1
mean_loss_index = mean_loss_index%30
return tuple([np.mean(loss) for loss in mean_loss_losses])
real_labels = -np.ones((batch_size,1))
fake_labels = np.ones((batch_size,1))
def train(input_img, vae_img):
critic_updates = 4
disc_real_loss = 0
disc_gen_loss = 0
for i in range(critic_updates):
fake_img = generator.predict(vae_img)
disc_real_loss += discriminator.train_on_batch([input_img, vae_img], real_labels)
disc_gen_loss += discriminator.train_on_batch([fake_img, vae_img], fake_labels)
disc_real_loss /= critic_updates
disc_gen_loss /= critic_updates
disc_total_loss = disc_real_loss + disc_gen_loss
gen_loss, _ = gan.train_on_batch([vae_img], real_labels)
rec_loss = 0
print('disc_total[%.6f] disc_real[%.6f] disc_gen[%.6f] gen_total[%.6f] reconstruction[%.3f]' % get_mean_loss((disc_total_loss, disc_real_loss, disc_gen_loss, gen_loss, rec_loss)))
print()
#Main training loop
def fit(train_ds):
global cp_num
while True:
for i, (data, _) in enumerate(train_ds.__iter__()):
print("Batch " + str(i) + '.')
input_image, vae_image, labels = data
input_image = tf.cast(input_image, dtype='float32')
vae_image = tf.cast(vae_image, dtype='float32')
if (i+1) % image_save_freq == 0:
saveImg()
if (i+1) % model_save_freq == 0:
print()
checkpoint.save(file_prefix = checkpoint_prefix)
print("Model Checkpoint <" + str(cp_num) + "> Saved!")
cp_num+=1
else:
print(str(model_save_freq - i%model_save_freq) + " batches until next save")
train(input_image, vae_image)
print()
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
fit(utk_data_gen(False))
Any help with this would be appreciated, even just to say that I haven't made any obvious errors and this issue may resolve itself with more training. I would have waited longer to see before posting the question, but I am under a time constraint to get something visual that shows progress.

Answer 1

I managed to fix this somewhat by doubling the number of filters in the critic, but also, interestingly, the results don't seem to be that much better than those I obtained by just adding a small amount of mae loss to the generator during testing. If anyone else runs into this problem I'd recommend considering this as a fix.