<jupyter_start><jupyter_text><jupyter_code>import tensorflow.keras
from tensorflow.keras import layers
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
import numpy as np

img_shape=(28, 28, 1)
batch_size = 16
latent_dim= 2

input_img = tensorflow.keras.Input(shape = img_shape)

x = layers.Conv2D(32,3,padding='same',activation='relu')(input_img)
x = layers.Conv2D(64,3,padding='same',activation='relu',strides=(2,2))(x)
x = layers.Conv2D(64,3,padding='same',activation='relu')(x)
x = layers.Conv2D(64,3,padding='same',activation='relu')(x)
shape_before_flattening = K.int_shape(x)

x = layers.Flatten()(x)
x = layers.Dense(32, activation='relu')(x)

z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)


def sampling(args):
  z_mean, z_log_var = args
  epsilon = K.random_normal(shape=(K.shape(z_mean)[0],latent_dim),mean=0.,stddev=1.)
  return z_mean+K.exp(z_log_var)* epsilon

z = layers.Lambda(sampling)([z_mean,z_log_var])

decoder_input = layers.Input(K.int_shape(z)[1:])
x = layers.Dense(np.prod(shape_before_flattening[1:]),activation='relu')(decoder_input)
x = layers.Reshape(shape_before_flattening[1:])(x)
x = layers.Conv2DTranspose(32, 3,padding='same',activation='relu',strides=(2,2))(x)
x = layers.Conv2D(1, 3, padding='same',activation='sigmoid')(x)

decoder = Model(decoder_input, x)
z_decoded = decoder(z)
class CustomVariationalLayer(tensorflow.keras.layers.Layer):

  def vae_loss(self, x, z_decoded):
    x = K.flatten(x)
    z_decoded = K.flatten(z_decoded)
    xent_loss = tensorflow.keras.metrics.binary_crossentropy(x, z_decoded)
    kl_loss = -5e-4 * K.mean(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
    return K.mean(xent_loss + kl_loss)

  def call(self, inputs):
    x = inputs[0]
    z_decoded = inputs[1]
    loss = self.vae_loss(x, z_decoded)
    self.add_loss(loss, inputs = inputs)
    return x

y = CustomVariationalLayer()([input_img, z_decoded])


from tensorflow.keras.datasets import mnist
vae=Model(input_img, y)
vae.compile(optimizer='rmsprop',loss=None)

(x_train, _), (x_test, y_test) = mnist.load_data()
x_train = x_train.astype('float32')/255.<jupyter_output>W0820 13:06:38.966424 140294730282880 training_utils.py:1101] Output custom_variational_layer missing from loss dictionary. We assume this was done on purpose. The fit and evaluate APIs will not be expecting any data to be passed to custom_variational_layer.
{'fraction_non_alphanumeric': 0.14619164619164618, 'fraction_numerical': 0.03644553644553645, 'mean_word_length': 6.403030303030303, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '510627', 'license': 'MIT', 'n_tokens_mistral': 1018, 'n_tokens_neox': 955, 'n_words': 165, 'path': 'VAE.ipynb', 'repo_name': 'fongabc/dcgan_code', 'size': '2442'}