Categories
Misc

VAE won’t learn color

I found the CVAE tutorial on tensorflow.org. I then repurposed it on new data however it is only reconstructing the outline in a nearly grayscale image, ignoring the colour, the loss is 14000. Please help…

#below: model

class VAE(tf.keras.Model):

def __init__(self, latent_dim):

super(VAE, self).__init__()

self.latent_dim = latent_dim

self.encoder = tf.keras.Sequential(

[

tf.keras.layers.InputLayer(input_shape=(128, 128,3)),

tf.keras.layers.Conv2D(

filters=32, kernel_size=3, strides=(2, 2), activation=’relu’),

tf.keras.layers.Conv2D(

filters=64, kernel_size=3, strides=(2, 2), activation=’relu’),

tf.keras.layers.Conv2D(

filters=128, kernel_size=3, strides=(2, 2), activation=’relu’),

# tf.keras.layers.Conv2D(

# filters=256, kernel_size=3, strides=(2, 2), activation=’relu’),

tf.keras.layers.Flatten(),

# No activation

tf.keras.layers.Dense(latent_dim + latent_dim),

]

)

self.decoder = tf.keras.Sequential(

[

tf.keras.layers.InputLayer(input_shape=(latent_dim,)),

tf.keras.layers.Dense(units=16*16*256, activation=tf.nn.relu),

tf.keras.layers.Reshape(target_shape=(16, 16, 256)),

tf.keras.layers.Conv2DTranspose(

filters=128, kernel_size=3, strides=2, padding=’same’,

activation=’relu’),

tf.keras.layers.Conv2DTranspose(

filters=64, kernel_size=3, strides=2, padding=’same’,

activation=’relu’),

tf.keras.layers.Conv2DTranspose(

filters=32, kernel_size=3, strides=2, padding=’same’,

activation=’relu’),

# No activation

tf.keras.layers.Conv2DTranspose(

filters=3, kernel_size=3, strides=1, padding=’same’),

]

)

u/tf.function

def sample(self, eps=None):

if eps is None:

eps = tf.random.normal(shape=(100, self.latent_dim))

return self.decode(eps, apply_sigmoid=True)

def encode(self, x):

mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)

return mean, logvar

def reparameterize(self, mean, logvar):

eps = tf.random.normal(shape=mean.shape)

return eps * tf.exp(logvar * .5) + mean

def decode(self, z, apply_sigmoid=False):

logits = self.decoder(z)

if apply_sigmoid:

probs = tf.sigmoid(logits)

return probs

return logits

optimizer = tf.keras.optimizers.Adam(1e-3)

def log_normal_pdf(sample, mean, logvar, raxis=1):

log2pi = tf.math.log(2. * np.pi)

return tf.reduce_sum(

-.5 * ((sample – mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),

axis=raxis)

def compute_loss(model, x):

mean, logvar = model.encode(x)

z = model.reparameterize(mean, logvar)

x_logit = model.decode(z)

cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)

logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])

logpz = log_normal_pdf(z, 0., 0.)

logqz_x = log_normal_pdf(z, mean, logvar)

return -tf.reduce_mean(logpx_z + logpz – logqz_x)

#above: loss

def train_step(model, x, optimizer):

“””Executes one training step and returns the loss.

This function computes the loss and gradients, and uses the latter to

update the model’s parameters.

“””

with tf.GradientTape() as tape:

loss = compute_loss(model, x)

gradients = tape.gradient(loss, model.trainable_variables)

optimizer.apply_gradients(zip(gradients, model.trainable_variables))

epochs = 10

latent_dim = 64 # up or 8

num_examples_to_generate = 32

model = VAE(latent_dim)

submitted by /u/much_bad_gramer
[visit reddit] [comments]

Leave a Reply

Your email address will not be published. Required fields are marked *