Initial v1

Author: aditya1601

infer.ipynb

@@ -0,0 +1,35 @@
+import tensorflow as tf
+from tensorflow import keras
+import tensorflow_datasets as tfds
+from vae import VarAutoEncoder
+
+
+# Load Data
+def preprocess(image, _):
+    """Return normalized image for both input and output label"""
+    image = tf.cast(image, tf.float32)/255.
+    image = tf.where(image < 0.5, 0., 1.)
+    return image, image
+
+
+ds_train, ds_test = tfds.load(
+    'mnist',
+    split=['train', 'test'],
+    as_supervised=True
+)
+ds_train = ds_train.map(preprocess).shuffle(1024).batch(64)
+ds_test = ds_test.map(preprocess).shuffle(1024).batch(64)
+
+
+# Train model
+model = VarAutoEncoder(latent_dim=2)
+model.compile(optimizer=keras.optimizers.Adam(1e-4))
+history = model.fit(
+    ds_train,
+    validation_data=ds_test,
+    epochs=50
+)
+
+
+# Save model
+model.save("models/vae-v1")

models/vae-v1/saved_model.pb

@@ -0,0 +1,83 @@
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras.layers import Conv2DTranspose, Conv2D, Dense, Flatten, Reshape
+import numpy as np
+
+
+class Sampling(keras.layers.Layer):
+    """Sample *z* from the *z_mean* and *z_logvar* from encoder to input in decoder"""
+
+    def call(self, inputs, **kwargs):
+        z_mean, z_logvar = inputs
+        epsilon = tf.random.normal(shape=tf.shape(z_mean))
+        return z_mean + tf.exp(0.5 * z_logvar) * epsilon
+
+
+class Encoder(keras.layers.Layer):
+    """Maps MNIST digits to triplet (z_mean, z_logvar, z)"""
+
+    def __init__(self, latent_dim, **kwargs):
+        super(Encoder, self).__init__(**kwargs)
+        self.conv1 = Conv2D(32, 3, (2,2), activation='relu')
+        self.conv2 = Conv2D(64, 3, (2,2), activation='relu')
+        self.flatten = Flatten()
+        self.dense3_1 = Dense(latent_dim)
+        self.dense3_2 = Dense(latent_dim)
+        self.sampling = Sampling()
+
+    def call(self, inputs, **kwargs):
+        x = self.conv1(inputs)
+        x = self.conv2(x)
+        x = self.flatten(x)
+        z_mean = self.dense3_1(x)
+        z_logvar = self.dense3_2(x)
+        z = self.sampling((z_mean, z_logvar))
+        return z_mean, z_logvar, z
+
+
+class Decoder(keras.layers.Layer):
+    """Reconstructs the image from latent variable *z*"""
+
+    def __init__(self, **kwargs):
+        super(Decoder, self).__init__(**kwargs)
+        self.dense1 = Dense(7*7*32, activation='relu')
+        self.reshape = Reshape((7, 7, 32))
+        self.deconv1 = Conv2DTranspose(64, 3, 2, padding='same', activation='relu')
+        self.deconv2 = Conv2DTranspose(32, 3, 2, padding='same', activation='relu')
+        self.out = Conv2DTranspose(1, 3, 1, padding='same')
+
+    def call(self, inputs, **kwargs):
+        x = self.dense1(inputs)
+        x = self.reshape(x)
+        x = self.deconv1(x)
+        x = self.deconv2(x)
+        return self.out(x)
+
+
+class VarAutoEncoder(keras.Model):
+    """Convolutional Variational AutoEncoder Model for MNIST"""
+
+    def __init__(self, latent_dim, **kwargs):
+        super(VarAutoEncoder, self).__init__(**kwargs)
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+    # TODO : Use TFP library functions
+    @tf.function
+    def log_normal_pdf(self, z, mean, logvar):
+        log2pi = tf.math.log(2. * np.pi)
+        return tf.reduce_sum(
+            -.5 * ((z - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
+            axis=1)
+
+    def call(self, inputs, **kwargs):
+        z_mean, z_logvar, z = self.encoder(inputs)
+        reconstructed = self.decoder(z)
+
+        # Compute loss
+        cross_entropy_loss = -tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(inputs, reconstructed), axis=[1, 2, 3])
+        kl_loss = self.log_normal_pdf(z, z_mean, z_logvar) - self.log_normal_pdf(z, 0., 0.)
+        total_loss = -tf.reduce_mean(cross_entropy_loss - kl_loss)
+
+        self.add_loss(total_loss)
+        return reconstructed