mercredi 5 février 2020

Learning material / Help on VAE with two Inference Nets / Encoders?

I'm really interested in this paper which uses VAEs to colourise black and white images https://arxiv.org/abs/1612.01958. I'd like to recreate the model, however, as I'm still new to developing Deep Learning models I've found it quite difficult to understand the code associated with it (https://github.com/aditya12agd5/divcolor).

I started from scratch with the basics and I've been able to make a VAE with one encoder and one decoder which can encode and decode colour images. Here's the code:

# Define Encoder
x_in = Input(shape=(train_uv.shape[1], train_uv.shape[2], train_uv.shape[3]))
layer = x_in
layer = Conv2D(filters=8, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=16, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=64, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
shape = K.int_shape(layer) # Shape before flattening

# Latent space
layer = Flatten()(layer)
layer = Dense(16, activation='relu')(layer)
z_mean = Dense(latent_dim)(layer)
z_log_var = Dense(latent_dim)(layer)
z = Lambda(sample, output_shape=(latent_dim,))([z_mean, z_log_var]) # Data passable to the decoder

# Define Decoder
latent_inputs = Input(shape=(latent_dim,))
layer = Dense(shape[1] * shape[2] * shape[3], activation='relu')(latent_inputs)
layer = Reshape((shape[1], shape[2], shape[3]))(layer)
layer = Conv2DTranspose(filters=64, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=16, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=8, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
# outputs = Conv2DTranspose(filters=1, kernel_size=kernel_size, activation='sigmoid', padding='same')(layer) # MNIST version
outputs = Conv2DTranspose(filters=2, kernel_size=kernel_size, activation='sigmoid', padding='same')(layer)

# VAE
vae = Model(x_in, x_out)

# Define loss
xent_loss = K.sum(K.binary_crossentropy(x_in, x_out), axis=[1, 2, 3])
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
vae_loss = K.mean(xent_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='rmsprop')

To make it then be able to colourise I understand that I need to add a second encoder which can take a grayscale image and use it to sample from the latent space to pull out a corresponding colour image from the decoder. I'm lost as to how I implement this separate encoder into the model though and I've been unable to find any learning material online for this type of VAE.

Does anyone know of any good tutorials or some readable Keras sample code, or could even give me an explanation as to how to implement a VAE which has two encoders?




Aucun commentaire:

Enregistrer un commentaire