Generative AI - Enhancing the VAE

Enhancing the VAE

Adding Convolutional Layers and Experimenting with Different Latent Space Dimensions

To make the VAE better, we can play around with different latent space variables and add more convolutional layers to the encoder and decoder.

1. Enhanced Encoder with Additional Convolutional Layers

def build_enhanced_encoder(latent_dim):
    inputs = layers.Input(shape=(28, 28, 1))
    x = layers.Conv2D(32, 3, activation='relu', strides=2, padding='same')(inputs)
    x = layers.Conv2D(64, 3, activation='relu', strides=2, padding='same')(x)
    x = layers.Conv2D(128, 3, activation='relu', strides=2, padding='same')(x)
    x = layers.Flatten()(x)
    x = layers.Dense(16, activation='relu')(x)
    z_mean = layers.Dense(latent_dim, name='z_mean')(x)
    z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)
    z = Sampling()([z_mean, z_log_var])
    return models.Model(inputs, [z_mean, z_log_var, z], name='encoder')

Additional Layer: To get more complicated features from the pictures that are fed in, an extra convolutional layer with 128 filters is added to the decoder. This makes it easier to learn a more complete hidden version.

2. Enhanced Decoder with Additional Convolutional Layers

def build_enhanced_decoder(latent_dim):
    latent_inputs = layers.Input(shape=(latent_dim,))
    x = layers.Dense(7 * 7 * 128, activation='relu')(latent_inputs)
    x = layers.Reshape((7, 7, 128))(x)
    x = layers.Conv2DTranspose(128, 3, strides=2, padding='same', activation='relu')(x)
    x = layers.Conv2DTranspose(64, 3, strides=2, padding='same', activation='relu')(x)
    x = layers.Conv2DTranspose(32, 3, strides=2, padding='same', activation='relu')(x)
    outputs = layers.Conv2DTranspose(1, 3, padding='same', activation='sigmoid')(x)
    return models.Model(latent_inputs, outputs, name='decoder')

Additional Layers: To improve the reconstruction of the input images, extra inverted convolutional layers with 128 and 64 filters are added to the decoder. These levels help improve the quality and organization of the picture.

3. Experimenting with Different Latent Dimensions

latent_dims = [2, 10, 50]

for latent_dim in latent_dims:
    encoder = build_enhanced_encoder(latent_dim)
    decoder = build_enhanced_decoder(latent_dim)
    vae = VAE(encoder, decoder)
    vae.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
    vae.fit(x_train, x_train, epochs=20, batch_size=64, validation_data=(x_test, x_test))
    print(f"Training complete for latent dimension: {latent_dim}")

Latent Dimensions: To see how the model works, we try making the latent space different sizes (2, 10, and 50 bits). It may be necessary to use more data and computing power to record more complicated patterns in hidden areas that are bigger.

By doing these steps, you can build and improve a VAE, learning more about how it works and how to make it run better by using convolutional layers and various latent space dimensions. This real code practice not only helps you understand VAEs better, but it also shows how flexible they are when dealing with complex data sets.