auto_encoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNXJm8flYkz9Itr7wjECns/",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/worthlessFella/deep_learning/blob/main/auto_encoder.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "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\n"
      ],
      "metadata": {
        "id": "0_84zLaULwBV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "vdE2rXJWLrsC"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "latent_dim = 2\n",
        "\n",
        "encoder_inputs = keras.Input(shape=(28, 28, 1))\n",
        "x = layers.Conv2D(32, 3, activation=\"relu\", strides=2, padding=\"same\")(encoder_inputs)\n",
        "x = layers.Conv2D(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
        "x = layers.Flatten()(x)\n",
        "x = layers.Dense(16, activation=\"relu\")(x)\n",
        "z_mean = layers.Dense(latent_dim, name=\"z_mean\")(x)\n",
        "z_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(x)\n",
        "z = Sampling()([z_mean, z_log_var])\n",
        "encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name=\"encoder\")\n",
        "encoder.summary()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TDM2vxDYMN6X",
        "outputId": "9a3fb4ef-7fe6-4f94-b7c6-84b3b5fac6db"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"encoder\"\n",
            "__________________________________________________________________________________________________\n",
            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            " input_1 (InputLayer)           [(None, 28, 28, 1)]  0           []                               \n",
            "                                                                                                  \n",
            " conv2d (Conv2D)                (None, 14, 14, 32)   320         ['input_1[0][0]']                \n",
            "                                                                                                  \n",
            " conv2d_1 (Conv2D)              (None, 7, 7, 64)     18496       ['conv2d[0][0]']                 \n",
            "                                                                                                  \n",
            " flatten (Flatten)              (None, 3136)         0           ['conv2d_1[0][0]']               \n",
            "                                                                                                  \n",
            " dense (Dense)                  (None, 16)           50192       ['flatten[0][0]']                \n",
            "                                                                                                  \n",
            " z_mean (Dense)                 (None, 2)            34          ['dense[0][0]']                  \n",
            "                                                                                                  \n",
            " z_log_var (Dense)              (None, 2)            34          ['dense[0][0]']                  \n",
            "                                                                                                  \n",
            " sampling (Sampling)            (None, 2)            0           ['z_mean[0][0]',                 \n",
            "                                                                  'z_log_var[0][0]']              \n",
            "                                                                                                  \n",
            "==================================================================================================\n",
            "Total params: 69,076\n",
            "Trainable params: 69,076\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "latent_inputs = keras.Input(shape=(latent_dim,))\n",
        "x = layers.Dense(7 * 7 * 64, activation=\"relu\")(latent_inputs)\n",
        "x = layers.Reshape((7, 7, 64))(x)\n",
        "x = layers.Conv2DTranspose(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
        "x = layers.Conv2DTranspose(32, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
        "decoder_outputs = layers.Conv2DTranspose(1, 3, activation=\"sigmoid\", padding=\"same\")(x)\n",
        "decoder = keras.Model(latent_inputs, decoder_outputs, name=\"decoder\")\n",
        "decoder.summary()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hju0C_fyNUVP",
        "outputId": "e271d428-db95-45f3-8398-b8557157dc77"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"decoder\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " input_2 (InputLayer)        [(None, 2)]               0         \n",
            "                                                                 \n",
            " dense_1 (Dense)             (None, 3136)              9408      \n",
            "                                                                 \n",
            " reshape (Reshape)           (None, 7, 7, 64)          0         \n",
            "                                                                 \n",
            " conv2d_transpose (Conv2DTra  (None, 14, 14, 64)       36928     \n",
            " nspose)                                                         \n",
            "                                                                 \n",
            " conv2d_transpose_1 (Conv2DT  (None, 28, 28, 32)       18464     \n",
            " ranspose)                                                       \n",
            "                                                                 \n",
            " conv2d_transpose_2 (Conv2DT  (None, 28, 28, 1)        289       \n",
            " ranspose)                                                       \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 65,089\n",
            "Trainable params: 65,089\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ]
        }
      ]
    },
    {
      "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",
        "                )\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": {
        "id": "m10ZlfHuN1M8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()\n",
        "mnist_digits = np.concatenate([x_train, x_test], axis=0)\n",
        "mnist_digits = np.expand_dims(mnist_digits, -1).astype(\"float32\") / 255\n",
        "\n",
        "vae = VAE(encoder, decoder)\n",
        "vae.compile(optimizer=keras.optimizers.Adam())\n",
        "vae.fit(mnist_digits, epochs=30, batch_size=128)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lhNj-RVbN6YF",
        "outputId": "bc30b10e-202c-49b7-f293-811955ffa513"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
            "11493376/11490434 [==============================] - 0s 0us/step\n",
            "11501568/11490434 [==============================] - 0s 0us/step\n",
            "Epoch 1/30\n",
            "547/547 [==============================] - 96s 171ms/step - loss: 256.6487 - reconstruction_loss: 207.9704 - kl_loss: 3.0486\n",
            "Epoch 2/30\n",
            "547/547 [==============================] - 94s 172ms/step - loss: 173.4493 - reconstruction_loss: 164.8763 - kl_loss: 5.7266\n",
            "Epoch 3/30\n",
            "547/547 [==============================] - 99s 181ms/step - loss: 165.0085 - reconstruction_loss: 158.1651 - kl_loss: 5.9590\n",
            "Epoch 4/30\n",
            "547/547 [==============================] - 94s 172ms/step - loss: 161.6189 - reconstruction_loss: 155.0986 - kl_loss: 6.1570\n",
            "Epoch 5/30\n",
            "547/547 [==============================] - 93s 170ms/step - loss: 159.9992 - reconstruction_loss: 153.2736 - kl_loss: 6.2132\n",
            "Epoch 6/30\n",
            "547/547 [==============================] - 91s 166ms/step - loss: 158.3538 - reconstruction_loss: 151.9497 - kl_loss: 6.2884\n",
            "Epoch 7/30\n",
            "547/547 [==============================] - 90s 165ms/step - loss: 157.1179 - reconstruction_loss: 150.8014 - kl_loss: 6.3342\n",
            "Epoch 8/30\n",
            "547/547 [==============================] - 91s 166ms/step - loss: 156.4142 - reconstruction_loss: 149.9175 - kl_loss: 6.3711\n",
            "Epoch 9/30\n",
            "547/547 [==============================] - 92s 167ms/step - loss: 155.4726 - reconstruction_loss: 149.1687 - kl_loss: 6.4017\n",
            "Epoch 10/30\n",
            "547/547 [==============================] - 91s 167ms/step - loss: 155.2036 - reconstruction_loss: 148.5273 - kl_loss: 6.4284\n",
            "Epoch 11/30\n",
            "547/547 [==============================] - 91s 167ms/step - loss: 154.4888 - reconstruction_loss: 148.0054 - kl_loss: 6.4367\n",
            "Epoch 12/30\n",
            "547/547 [==============================] - 91s 167ms/step - loss: 153.8256 - reconstruction_loss: 147.5280 - kl_loss: 6.4377\n",
            "Epoch 13/30\n",
            "547/547 [==============================] - 91s 167ms/step - loss: 153.6631 - reconstruction_loss: 147.1826 - kl_loss: 6.4629\n",
            "Epoch 14/30\n",
            "547/547 [==============================] - 91s 167ms/step - loss: 153.5042 - reconstruction_loss: 146.8392 - kl_loss: 6.4312\n",
            "Epoch 15/30\n",
            "547/547 [==============================] - 92s 168ms/step - loss: 153.2484 - reconstruction_loss: 146.5130 - kl_loss: 6.4414\n",
            "Epoch 16/30\n",
            "547/547 [==============================] - 92s 168ms/step - loss: 152.9648 - reconstruction_loss: 146.1715 - kl_loss: 6.4317\n",
            "Epoch 17/30\n",
            "547/547 [==============================] - 92s 167ms/step - loss: 152.4683 - reconstruction_loss: 146.0751 - kl_loss: 6.4508\n",
            "Epoch 18/30\n",
            "547/547 [==============================] - 94s 172ms/step - loss: 152.8564 - reconstruction_loss: 145.7555 - kl_loss: 6.4289\n",
            "Epoch 19/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 151.7261 - reconstruction_loss: 145.5540 - kl_loss: 6.4311\n",
            "Epoch 20/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 152.2215 - reconstruction_loss: 145.3264 - kl_loss: 6.4365\n",
            "Epoch 21/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 151.6183 - reconstruction_loss: 145.0876 - kl_loss: 6.4332\n",
            "Epoch 22/30\n",
            "547/547 [==============================] - 95s 174ms/step - loss: 151.0392 - reconstruction_loss: 144.9400 - kl_loss: 6.4366\n",
            "Epoch 23/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 151.2059 - reconstruction_loss: 144.8373 - kl_loss: 6.4236\n",
            "Epoch 24/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 151.0601 - reconstruction_loss: 144.6231 - kl_loss: 6.4384\n",
            "Epoch 25/30\n",
            "547/547 [==============================] - 94s 172ms/step - loss: 150.6392 - reconstruction_loss: 144.4651 - kl_loss: 6.4425\n",
            "Epoch 26/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 150.7789 - reconstruction_loss: 144.3312 - kl_loss: 6.4186\n",
            "Epoch 27/30\n",
            "547/547 [==============================] - 95s 174ms/step - loss: 150.6334 - reconstruction_loss: 144.2263 - kl_loss: 6.4377\n",
            "Epoch 28/30\n",
            "547/547 [==============================] - 95s 173ms/step - loss: 150.4016 - reconstruction_loss: 144.0788 - kl_loss: 6.4395\n",
            "Epoch 29/30\n",
            "547/547 [==============================] - 95s 174ms/step - loss: 150.1742 - reconstruction_loss: 143.9644 - kl_loss: 6.4371\n",
            "Epoch 30/30\n",
            "547/547 [==============================] - 95s 174ms/step - loss: 150.2537 - reconstruction_loss: 143.8587 - kl_loss: 6.4436\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7f18a5995a90>"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    }
  ]
}