{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\ntf.config.list_physical_devices('GPU')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Sampling(layers.Layer):\n    \"\"\"Uses (z_mean, z_log_var) to sample z, the vector encoding a digit.\"\"\"\n\n    def call(self, inputs):\n        z_mean, z_log_var = inputs\n        batch = tf.shape(z_mean)[0]\n        dim = tf.shape(z_mean)[1]\n        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n        return z_mean + tf.exp(0.5 * z_log_var) * epsilon","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"latent_dim = 100\n\nencoder_inputs = keras.Input(shape=(28, 28, 1))\nx = layers.Conv2D(32, 3, activation=\"relu\", strides=2, padding=\"same\")(encoder_inputs)\nx = layers.Conv2D(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\nx = layers.Flatten()(x)\nx = layers.Dense(160, activation=\"relu\")(x)\nz_mean = layers.Dense(latent_dim, name=\"z_mean\")(x)\nz_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(x)\nz = Sampling()([z_mean, z_log_var])\nencoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name=\"encoder\")\nencoder.summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"latent_inputs = keras.Input(shape=(latent_dim,))\nx = layers.Dense(7 * 7 * 64, activation=\"relu\")(latent_inputs)\nx = layers.Reshape((7, 7, 64))(x)\nx = layers.Conv2DTranspose(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\nx = layers.Conv2DTranspose(32, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\ndecoder_outputs = layers.Conv2DTranspose(1, 3, activation=\"sigmoid\", padding=\"same\")(x)\ndecoder = keras.Model(latent_inputs, decoder_outputs, name=\"decoder\")\ndecoder.summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class VAE(keras.Model):\n    def __init__(self, encoder, decoder, **kwargs):\n        super(VAE, self).__init__(**kwargs)\n        self.encoder = encoder\n        self.decoder = decoder\n        self.total_loss_tracker = keras.metrics.Mean(name=\"total_loss\")\n        self.reconstruction_loss_tracker = keras.metrics.Mean(\n            name=\"reconstruction_loss\"\n        )\n        self.kl_loss_tracker = keras.metrics.Mean(name=\"kl_loss\")\n\n    @property\n    def metrics(self):\n        return [\n            self.total_loss_tracker,\n            self.reconstruction_loss_tracker,\n            self.kl_loss_tracker,\n        ]\n\n    def train_step(self, data):\n        with tf.GradientTape() as tape:\n            z_mean, z_log_var, z = self.encoder(data)\n            reconstruction = self.decoder(z)\n            reconstruction_loss = tf.reduce_mean(\n                tf.reduce_sum(\n#                     keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2)\n                    keras.losses.mse(data, reconstruction), axis=(1, 2)\n                )\n            )\n            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))\n            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))\n            total_loss = reconstruction_loss + kl_loss\n        grads = tape.gradient(total_loss, self.trainable_weights)\n        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n        self.total_loss_tracker.update_state(total_loss)\n        self.reconstruction_loss_tracker.update_state(reconstruction_loss)\n        self.kl_loss_tracker.update_state(kl_loss)\n        return {\n            \"loss\": self.total_loss_tracker.result(),\n            \"reconstruction_loss\": self.reconstruction_loss_tracker.result(),\n            \"kl_loss\": self.kl_loss_tracker.result(),\n        }\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\nmnist_digits = np.concatenate([x_train, x_test], axis=0)\nmnist_digits = np.expand_dims(mnist_digits, -1).astype(\"float32\") / 255\nmnist_digits.shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"directory = \"/kaggle/input/dog-breed-identification/\"\nds = keras.preprocessing.image_dataset_from_directory(\n    directory, labels='inferred', label_mode='int',\n    class_names=None, color_mode='grayscale', batch_size=70000, image_size=(28,\n    28), shuffle=True, seed=None, validation_split=None, subset=None,\n    interpolation='bilinear', follow_links=False\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"imagenet_gen = next( (X.numpy()/255 for X, y in ds) )\nimagenet_gen.shape\ndata = imagenet_gen","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = mnist_digits","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vae = VAE(encoder, decoder)\nvae.compile(optimizer=keras.optimizers.Adam())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = vae.fit(data, epochs=30, batch_size=128)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"z_mean, z_log_var, z = vae.encoder.predict(data[0:1,...])\ndec = vae.decoder.predict(z)\ndec2 = vae.decoder.predict(z_mean)\ndec3 = vae.decoder.predict(z_mean - 0.5*np.exp(z_log_var))\ndec4 = vae.decoder.predict(z_mean + 0.5*np.exp(z_log_var))\nprint(z_mean.mean())\nprint(z_log_var.std())\nprint(z.std())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 5, figsize=(10, 4))\nax[0].imshow(np.uint8(data[0,...]*255), cmap='gray')\nax[1].imshow(np.uint8(dec[0,...]*255), cmap='gray')\nax[2].imshow(np.uint8(dec2[0,...]*255), cmap='gray')\nax[3].imshow(np.uint8(dec3[0,...]*255), cmap='gray')\nax[4].imshow(np.uint8(dec4[0,...]*255), cmap='gray')\nfor ax_ in ax:\n    ax_.axis('off')\nfig.tight_layout()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(history.history['reconstruction_loss'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(history.history['kl_loss'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}