{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1eY-Qmf8ItWv"
      },
      "source": [
        "# Autoencoders"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yEqo3q8JI5uG"
      },
      "source": [
        "When training CNNs, one of the problems is that we need a lot of labeled data. In the case of image classification, we need to separate images into different classes, which is a manual effort.\n",
        "\n",
        "However, we might want to use raw (unlabeled) data for training CNN feature extractors, which is called **self-supervised learning**. Instead of labels, we will use training images as both network input and output. The main idea of **autoencoder** is that we will have an **encoder network** that converts input image into some **latent space** (normally it is just a vector of some smaller size), the then **decoder network**, whose goal would be to reconstruct the original image.\n",
        "\n",
        "Since we are training autoencoder to capture as much of the information from the original image as possible for accurate reconstruction, the network tries to find the best **embedding** of input images to capture the meaning.\n",
        "\n",
        "![AutoEncoder Diagram](images/autoencoder_schema.jpg)\n",
        "\n",
        "*Image from [Keras blog](https://blog.keras.io/building-autoencoders-in-keras.html)*\n",
        "\n",
        "Most of the examples below are inspired by [this article](https://blog.keras.io/building-autoencoders-in-keras.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZ38ZCGRYhHh"
      },
      "source": [
        "Let's create simplest autoencoder for MNIST:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "48ua_UfDIqum",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.datasets import mnist\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "(x_train, y_trainclass), (x_test, y_testclass) = mnist.load_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        },
        "id": "-VvA6_e_Ytv9",
        "outputId": "7f43ddd9-2044-4c88-a57d-7dffbb0741fe",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 17140 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "def plotn(n,x):\n",
        "  fig,ax = plt.subplots(1,n)\n",
        "  for i,z in enumerate(x[0:n]):\n",
        "    ax[i].imshow(z.reshape(28,28) if z.size==28*28 else z.reshape(14,14) if z.size==14*14 else z)\n",
        "  plt.show()\n",
        "  \n",
        "plotn(5,x_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "XVXKfTiKZGzs",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Lambda\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.losses import binary_crossentropy,mse\n",
        "\n",
        "input_img = Input(shape=(28, 28, 1))\n",
        "\n",
        "x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)\n",
        "x = MaxPooling2D((2, 2), padding='same')(x)\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)\n",
        "x = MaxPooling2D((2, 2), padding='same')(x)\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)\n",
        "encoded = MaxPooling2D((2, 2), padding='same')(x)\n",
        "\n",
        "encoder = Model(input_img,encoded)\n",
        "\n",
        "input_rep = Input(shape=(4,4,8))\n",
        "\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(input_rep)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "x = Conv2D(16, (3, 3), activation='relu')(x)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)\n",
        "\n",
        "decoder = Model(input_rep,decoded)\n",
        "\n",
        "autoencoder = Model(input_img, decoder(encoder(input_img)))\n",
        "autoencoder.compile(optimizer='adam', loss='binary_crossentropy')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "xRVImJGBZ2kt",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "x_train = x_train.astype('float32') / 255.\n",
        "x_test = x_test.astype('float32') / 255.\n",
        "x_train = np.reshape(x_train, (len(x_train), 28, 28, 1))\n",
        "x_test = np.reshape(x_test, (len(x_test), 28, 28, 1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 131,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H6YD611WaAAG",
        "outputId": "3cf9f53c-8898-4bde-89b7-3670804c2373",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train on 60000 samples, validate on 10000 samples\n",
            "Epoch 1/25\n",
            "59648/60000 [============================>.] - ETA: 0s - loss: 0.2134"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "60000/60000 [==============================] - 6s 99us/sample - loss: 0.2130 - val_loss: 0.1454\n",
            "Epoch 2/25\n",
            "60000/60000 [==============================] - 5s 86us/sample - loss: 0.1353 - val_loss: 0.1258\n",
            "Epoch 3/25\n",
            "60000/60000 [==============================] - 5s 87us/sample - loss: 0.1225 - val_loss: 0.1177\n",
            "Epoch 4/25\n",
            "60000/60000 [==============================] - 5s 85us/sample - loss: 0.1163 - val_loss: 0.1126\n",
            "Epoch 5/25\n",
            "60000/60000 [==============================] - 5s 87us/sample - loss: 0.1120 - val_loss: 0.1091\n",
            "Epoch 6/25\n",
            "60000/60000 [==============================] - 5s 86us/sample - loss: 0.1093 - val_loss: 0.1070\n",
            "Epoch 7/25\n",
            "60000/60000 [==============================] - 5s 87us/sample - loss: 0.1072 - val_loss: 0.1055\n",
            "Epoch 8/25\n",
            "60000/60000 [==============================] - 5s 87us/sample - loss: 0.1057 - val_loss: 0.1041\n",
            "Epoch 9/25\n",
            "60000/60000 [==============================] - 5s 85us/sample - loss: 0.1045 - val_loss: 0.1028\n",
            "Epoch 10/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.1035 - val_loss: 0.1022\n",
            "Epoch 11/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.1026 - val_loss: 0.1011\n",
            "Epoch 12/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.1018 - val_loss: 0.1003\n",
            "Epoch 13/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.1012 - val_loss: 0.0996\n",
            "Epoch 14/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.1005 - val_loss: 0.0991\n",
            "Epoch 15/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.1000 - val_loss: 0.0988\n",
            "Epoch 16/25\n",
            "60000/60000 [==============================] - 5s 82us/sample - loss: 0.0995 - val_loss: 0.0981\n",
            "Epoch 17/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.0990 - val_loss: 0.0976\n",
            "Epoch 18/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.0986 - val_loss: 0.0974\n",
            "Epoch 19/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.0982 - val_loss: 0.0969\n",
            "Epoch 20/25\n",
            "60000/60000 [==============================] - 5s 85us/sample - loss: 0.0978 - val_loss: 0.0970\n",
            "Epoch 21/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.0975 - val_loss: 0.0962\n",
            "Epoch 22/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.0971 - val_loss: 0.0960\n",
            "Epoch 23/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.0968 - val_loss: 0.0958\n",
            "Epoch 24/25\n",
            "60000/60000 [==============================] - 5s 84us/sample - loss: 0.0966 - val_loss: 0.0953\n",
            "Epoch 25/25\n",
            "60000/60000 [==============================] - 5s 83us/sample - loss: 0.0963 - val_loss: 0.0953\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f3fa179b690>"
            ]
          },
          "execution_count": 131,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "autoencoder.fit(x_train, x_train,\n",
        "                epochs=25,\n",
        "                batch_size=128,\n",
        "                shuffle=True,\n",
        "                validation_data=(x_test, x_test))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 132,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 256
        },
        "id": "CTEBvy2FaRBF",
        "outputId": "b18d0d16-c88f-47bb-fe2e-2ba5871f0563",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 15840 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 18520 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "y_test = autoencoder.predict(x_test[0:5])\n",
        "plotn(5,x_test)\n",
        "plotn(5,y_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 133,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "44u_8PR7kaoY",
        "outputId": "71916b80-0821-4cd2-c481-bbbd5dac2fd4",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        }
      ],
      "source": [
        "encoder = Model(input_img, encoded)\n",
        "encoded_imgs = encoder.predict(x_test[0:5])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 134,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 164
        },
        "id": "DJS4ZdO5oOFn",
        "outputId": "89df7031-c58b-46ae-9fc7-4ddbf5e46ef9",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 9856 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "plotn(5,encoded_imgs.reshape(5,-1,8))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 136,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        },
        "id": "WXIys9NGz8bF",
        "outputId": "c0b2ede8-a041-4cac-ccf6-3e40817df8c5",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "6.3110805 0.0\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 5444 characters]",
            "text/plain": [
              "<Figure size 432x288 with 7 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "print(encoded_imgs.max(),encoded_imgs.min())\n",
        "res = decoder.predict(7*np.random.rand(7,4,4,8))\n",
        "plotn(7,res)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5RQtFfjpNuI6"
      },
      "source": [
        "> **Task 1**: Try to train autoencoder with very small latent vector size, eg. 2, and plot the dots corresponding to different digits. *Hint: Use fully-connected dense layer after the convoluitonal part to reduce the vector size to the required value.*\n",
        "\n",
        "> **Task 2**: Starting from different digits, obtain their latent space representations, and see what effect adding some noise to the latent space has on the resulting digits."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xMLzTpExtPl1"
      },
      "source": [
        "## Denoising\n",
        "\n",
        "Autoencoders can be effectively used to remove noise from images. In order to train denoiser, we will start with noise-free images, and add artificial noise to them. Then, we will feed autoencoder with noisy images as input, and noise-free images as output.\n",
        "\n",
        "Let's see how this works for MNIST:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 137,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        },
        "id": "uVprJu59lte0",
        "outputId": "3e9e924b-ec01-4ff3-c49a-3dd70db84db1",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 68504 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "def noisify(data):\n",
        "  return np.clip(data+np.random.normal(loc=0.5,scale=0.5,size=data.shape),0.,1.)\n",
        "\n",
        "x_train_noise = noisify(x_train)\n",
        "x_test_noise = noisify(x_test)\n",
        "\n",
        "plotn(5,x_train_noise)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 141,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mlwxbuaHomTr",
        "outputId": "24c6c3be-e755-4216-815a-46561448fe7f",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train on 60000 samples, validate on 10000 samples\n",
            "Epoch 1/25\n",
            "60000/60000 [==============================] - 6s 101us/sample - loss: 0.1576 - val_loss: 0.1566\n",
            "Epoch 2/25\n",
            "60000/60000 [==============================] - 6s 95us/sample - loss: 0.1564 - val_loss: 0.1553\n",
            "Epoch 3/25\n",
            "60000/60000 [==============================] - 6s 94us/sample - loss: 0.1555 - val_loss: 0.1539\n",
            "Epoch 4/25\n",
            "60000/60000 [==============================] - 6s 95us/sample - loss: 0.1545 - val_loss: 0.1530\n",
            "Epoch 5/25\n",
            "60000/60000 [==============================] - 6s 95us/sample - loss: 0.1538 - val_loss: 0.1517\n",
            "Epoch 6/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1528 - val_loss: 0.1506\n",
            "Epoch 7/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1521 - val_loss: 0.1499\n",
            "Epoch 8/25\n",
            "60000/60000 [==============================] - 5s 92us/sample - loss: 0.1514 - val_loss: 0.1495\n",
            "Epoch 9/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1508 - val_loss: 0.1487\n",
            "Epoch 10/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1500 - val_loss: 0.1483\n",
            "Epoch 11/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1495 - val_loss: 0.1484\n",
            "Epoch 12/25\n",
            "60000/60000 [==============================] - 6s 94us/sample - loss: 0.1487 - val_loss: 0.1468\n",
            "Epoch 13/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1482 - val_loss: 0.1467\n",
            "Epoch 14/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1476 - val_loss: 0.1459\n",
            "Epoch 15/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1469 - val_loss: 0.1450\n",
            "Epoch 16/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1463 - val_loss: 0.1442\n",
            "Epoch 17/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1457 - val_loss: 0.1441\n",
            "Epoch 18/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1451 - val_loss: 0.1429\n",
            "Epoch 19/25\n",
            "60000/60000 [==============================] - 6s 92us/sample - loss: 0.1445 - val_loss: 0.1425\n",
            "Epoch 20/25\n",
            "60000/60000 [==============================] - 6s 94us/sample - loss: 0.1440 - val_loss: 0.1418\n",
            "Epoch 21/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1435 - val_loss: 0.1423\n",
            "Epoch 22/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1430 - val_loss: 0.1409\n",
            "Epoch 23/25\n",
            "60000/60000 [==============================] - 6s 94us/sample - loss: 0.1426 - val_loss: 0.1405\n",
            "Epoch 24/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1422 - val_loss: 0.1409\n",
            "Epoch 25/25\n",
            "60000/60000 [==============================] - 6s 93us/sample - loss: 0.1418 - val_loss: 0.1398\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f3fa612c4d0>"
            ]
          },
          "execution_count": 141,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "autoencoder.fit(x_train_noise, x_train,\n",
        "                epochs=25,\n",
        "                batch_size=128,\n",
        "                shuffle=True,\n",
        "                validation_data=(x_test_noise, x_test))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 144,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "7kctWkcwv3kd",
        "outputId": "00bfe5d9-882b-46d3-b8b2-8310756bc9f9",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 68128 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 22140 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "y_test = autoencoder.predict(x_test_noise[0:5])\n",
        "plotn(5,x_test_noise)\n",
        "plotn(5,y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C_u1vG1BDqVx"
      },
      "source": [
        "> **Exercise:** See how denoiser trained on MNIST digits works for different images. As an example, you can take [Fashion MNIST](https://keras.io/api/datasets/fashion_mnist/) dataset, which has the same image size. Note that denoiser works well only on the same image type that is was trained on (i.e. for the same probability distribution of input data)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fcjt-egdU6Kw"
      },
      "source": [
        "## Super-resolution\n",
        "\n",
        "Similarly to denoiser, we can train autoencoders to increase the resolution of the image. To train super-resolution network, we will start with high-res images, and automatically downscale them to produce network inputs. We will then feed autoencoder with small images as inputs and high-res images as outputs.\n",
        "\n",
        "Let's downscale MNIST to 14x14:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 108
        },
        "id": "JGV724C9VQ7m",
        "outputId": "d63d8ad4-55bc-431b-bc95-79f994713ad3",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 4780 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "x_train_lr = tf.keras.layers.AveragePooling2D()(x_train).numpy()\n",
        "x_test_lr = tf.keras.layers.AveragePooling2D()(x_test).numpy()\n",
        "plotn(5,x_train_lr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "VXIFG2ulYuGM",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Lambda\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.losses import binary_crossentropy,mse\n",
        "\n",
        "input_img = Input(shape=(14, 14, 1))\n",
        "\n",
        "x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)\n",
        "x = MaxPooling2D((2, 2), padding='same')(x)\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)\n",
        "encoded = MaxPooling2D((2, 2), padding='same')(x)\n",
        "\n",
        "encoder = Model(input_img,encoded)\n",
        "\n",
        "input_rep = Input(shape=(4,4,8))\n",
        "\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(input_rep)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "x = Conv2D(16, (3, 3), activation='relu')(x)\n",
        "x = UpSampling2D((2, 2))(x)\n",
        "decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)\n",
        "\n",
        "decoder = Model(input_rep,decoded)\n",
        "\n",
        "autoencoder = Model(input_img, decoder(encoder(input_img)))\n",
        "autoencoder.compile(optimizer='adam', loss='binary_crossentropy')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tie6d1OCY2ql",
        "outputId": "0a97cea6-3fd3-41e3-c743-e332535fa779",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1/25\n",
            "469/469 [==============================] - 6s 10ms/step - loss: 0.3413 - val_loss: 0.1519\n",
            "Epoch 2/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1457 - val_loss: 0.1292\n",
            "Epoch 3/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1273 - val_loss: 0.1202\n",
            "Epoch 4/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1189 - val_loss: 0.1142\n",
            "Epoch 5/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1148 - val_loss: 0.1107\n",
            "Epoch 6/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1115 - val_loss: 0.1083\n",
            "Epoch 7/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1093 - val_loss: 0.1063\n",
            "Epoch 8/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1071 - val_loss: 0.1046\n",
            "Epoch 9/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1060 - val_loss: 0.1037\n",
            "Epoch 10/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1048 - val_loss: 0.1026\n",
            "Epoch 11/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1039 - val_loss: 0.1019\n",
            "Epoch 12/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1030 - val_loss: 0.1012\n",
            "Epoch 13/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1024 - val_loss: 0.1004\n",
            "Epoch 14/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1017 - val_loss: 0.0999\n",
            "Epoch 15/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1010 - val_loss: 0.0993\n",
            "Epoch 16/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.1005 - val_loss: 0.0989\n",
            "Epoch 17/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0999 - val_loss: 0.0983\n",
            "Epoch 18/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0995 - val_loss: 0.0982\n",
            "Epoch 19/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0990 - val_loss: 0.0975\n",
            "Epoch 20/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0987 - val_loss: 0.0971\n",
            "Epoch 21/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0981 - val_loss: 0.0971\n",
            "Epoch 22/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0979 - val_loss: 0.0965\n",
            "Epoch 23/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0977 - val_loss: 0.0959\n",
            "Epoch 24/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0972 - val_loss: 0.0957\n",
            "Epoch 25/25\n",
            "469/469 [==============================] - 4s 9ms/step - loss: 0.0972 - val_loss: 0.0955\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f66790ada90>"
            ]
          },
          "execution_count": 8,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "autoencoder.fit(x_train_lr, x_train,\n",
        "                epochs=25,\n",
        "                batch_size=128,\n",
        "                shuffle=True,\n",
        "                validation_data=(x_test_lr, x_test))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 200
        },
        "id": "C-wLl0BiZGR4",
        "outputId": "f0104552-e056-4cbf-a238-09820839d70f",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 4796 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 19188 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "y_test_lr = autoencoder.predict(x_test_lr[0:5])\n",
        "plotn(5,x_test_lr)\n",
        "plotn(5,y_test_lr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BlnJnDwfZ_f0"
      },
      "source": [
        "> **Exercise**: Try to train super-resolution network on [CIFAR-10](https://keras.io/api/datasets/cifar10/) for 2x and 4x upscaling. Use noise as input to 4x upscaling model and observe the result."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IYKb6NwH1YOL"
      },
      "source": [
        "## Variational Auto-Encoders (VAE)\n",
        "\n",
        "Traditional autoencoders reduce the dimension of the input data somehow, figuring out the important features of input images. However, latent vectors ofter do not make much sense. In other words, taking MNIST dataset as an example, figuring out which digits correspond to different latent vectors is not an easy task, because close latent vectors would not necessarily correspond to the same digits. \n",
        "\n",
        "On the other hand, to train *generative* models it is better to have some understanding of the latent space. This idea leads us to **variational autoencoder** (VAE).\n",
        "\n",
        "VAE is the autoencoder that learns to predict *statistical distribution* of the latent parameters, so-called **latent distribution**. For example, we can assume that latent vectors would be distributed as $N(\\mathrm{z\\_mean},e^{\\mathrm{z\\_log\\_sigma}})$, where $\\mathrm{z\\_mean}, \\mathrm{z\\_log\\_sigma} \\in\\mathbb{R}^d$. Encoder in VAE learns to predict those parameters, and then decoder takes a random vector from this distribution to reconstruct the object.\n",
        "\n",
        "To summarize:\n",
        "\n",
        " * From input vector, we predict `z_mean` and `z_log_sigma` (instead of predicting the standard deviation itself, we predict it's logarithm)\n",
        " * We sample a vector `sample` from the distribution $N(\\mathrm{z\\_mean},e^{\\mathrm{z\\_log\\_sigma}})$\n",
        " * Decoder tries to decode the original image using `sample` as an input vector\n",
        "\n",
        " <img src=\"images/vae.png\" width=\"50%\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "a8bqQpG6wuJ9",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "intermediate_dim = 512\n",
        "latent_dim = 2\n",
        "batch_size = 128\n",
        "\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "inputs = Input(shape=(784,))\n",
        "h = Dense(intermediate_dim, activation='relu')(inputs)\n",
        "z_mean = Dense(latent_dim)(h)\n",
        "z_log_sigma = Dense(latent_dim)(h)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "y6cJ93kwQkM3",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def sampling(args):\n",
        "    z_mean, z_log_sigma = args\n",
        "    bs = tf.shape(z_mean)[0]\n",
        "    epsilon = tf.random.normal(shape=(bs, latent_dim))\n",
        "    return z_mean + tf.exp(z_log_sigma) * epsilon\n",
        "\n",
        "z = Lambda(sampling)([z_mean, z_log_sigma])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "Bf3UqBMXQwpQ",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "encoder = Model(inputs, [z_mean, z_log_sigma, z])\n",
        "\n",
        "latent_inputs = Input(shape=(latent_dim,))\n",
        "x = Dense(intermediate_dim, activation='relu')(latent_inputs)\n",
        "outputs = Dense(784, activation='sigmoid')(x)\n",
        "\n",
        "decoder = Model(latent_inputs, outputs)\n",
        "\n",
        "outputs = decoder(encoder(inputs)[2])\n",
        "\n",
        "vae = Model(inputs, outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kY24IoGC_laa"
      },
      "source": [
        "Variational auto-encoders use complex loss function that consists of two parts:\n",
        "* **Reconstruction loss** is the loss function that shows how close reconstructed image is to the target (can be MSE). It is the same loss function as in normal autoencoders.\n",
        "* **KL loss**, which ensures that latent variable distributions stays close to normal distribution. It is based on the notion of [Kullback-Leibler divergence](https://www.countbayesie.com/blog/2017/5/9/kullback-leibler-divergence-explained) - a metric to estimate how similar two statistical distributions are."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "czXAjHUJR-BE",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def vae_loss(x1,x2):\n",
        "  reconstruction_loss = mse(x1,x2)*784\n",
        "  tmp = 1 + z_log_sigma - tf.square(z_mean) - tf.exp(z_log_sigma)\n",
        "  kl_loss = -0.5*tf.reduce_sum(tmp, axis=-1)\n",
        "  return tf.convert_to_tensor(tf.reduce_mean(reconstruction_loss + kl_loss))\n",
        "\n",
        "vae.compile(optimizer='rmsprop', loss=vae_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zky-60lERlG-",
        "outputId": "baed11e4-b165-42d5-d86b-1037741554c8",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train on 60000 samples, validate on 10000 samples\n",
            "Epoch 1/25\n",
            "59520/60000 [============================>.] - ETA: 0s - loss: 48.6396"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "60000/60000 [==============================] - 4s 64us/sample - loss: 48.5874 - val_loss: 41.8877\n",
            "Epoch 2/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 41.1296 - val_loss: 40.2556\n",
            "Epoch 3/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 40.0063 - val_loss: 39.3692\n",
            "Epoch 4/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 39.2531 - val_loss: 38.7666\n",
            "Epoch 5/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 38.7147 - val_loss: 38.6124\n",
            "Epoch 6/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 38.2962 - val_loss: 38.1867\n",
            "Epoch 7/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 37.9756 - val_loss: 37.9831\n",
            "Epoch 8/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 37.6933 - val_loss: 37.5475\n",
            "Epoch 9/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 37.4323 - val_loss: 37.2913\n",
            "Epoch 10/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 37.2133 - val_loss: 37.1992\n",
            "Epoch 11/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.9966 - val_loss: 36.9521\n",
            "Epoch 12/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.8204 - val_loss: 36.8431\n",
            "Epoch 13/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.6490 - val_loss: 36.6979\n",
            "Epoch 14/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.5023 - val_loss: 36.6661\n",
            "Epoch 15/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.3456 - val_loss: 36.4957\n",
            "Epoch 16/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 36.2266 - val_loss: 36.6669\n",
            "Epoch 17/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 36.1045 - val_loss: 36.4855\n",
            "Epoch 18/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 35.9922 - val_loss: 36.4150\n",
            "Epoch 19/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 35.8968 - val_loss: 36.1196\n",
            "Epoch 20/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 35.7991 - val_loss: 36.0708\n",
            "Epoch 21/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 35.7129 - val_loss: 36.1686\n",
            "Epoch 22/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 35.6214 - val_loss: 36.1080\n",
            "Epoch 23/25\n",
            "60000/60000 [==============================] - 3s 57us/sample - loss: 35.5357 - val_loss: 36.2309\n",
            "Epoch 24/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 35.4528 - val_loss: 36.1416\n",
            "Epoch 25/25\n",
            "60000/60000 [==============================] - 3s 56us/sample - loss: 35.3650 - val_loss: 35.7258\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f0e00233890>"
            ]
          },
          "execution_count": 25,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_train_flat = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))\n",
        "x_test_flat = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))\n",
        "\n",
        "vae.fit(x_train_flat, x_train_flat,\n",
        "        shuffle=True,\n",
        "        epochs=25,\n",
        "        batch_size=batch_size,\n",
        "        validation_data=(x_test_flat, x_test_flat))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 256
        },
        "id": "YF8hHAZVR5x_",
        "outputId": "146edef1-676d-4592-8b9f-250ff4b98fa6",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 15840 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 22448 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "y_test = vae.predict(x_test_flat[0:5])\n",
        "plotn(5,x_test_flat)\n",
        "plotn(5,y_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 433
        },
        "id": "jZ0PzYmInDi5",
        "outputId": "b114fa04-ba23-47e3-b63e-1e3eac248463",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 86864 characters]",
            "text/plain": [
              "<Figure size 432x432 with 2 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "x_test_encoded = encoder.predict(x_test_flat)[0]\n",
        "plt.figure(figsize=(6, 6))\n",
        "plt.scatter(x_test_encoded[:, 0], x_test_encoded[:, 1], c=y_testclass)\n",
        "plt.colorbar()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 303
        },
        "id": "AZQfTelyOigw",
        "outputId": "3c0fc864-e3e5-4b57-d84f-572358f2e61b",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
            "  warnings.warn('`Model.state_updates` will be removed in a future version. '\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 72492 characters]",
            "text/plain": [
              "<Figure size 432x288 with 100 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "def plotsample(n):\n",
        "  dx = np.linspace(-1,1,n)\n",
        "  dy = np.linspace(-1,1,n)\n",
        "  fig,ax = plt.subplots(n,n)\n",
        "  for i,xi in enumerate(dx):\n",
        "    for j,xj in enumerate(dy):\n",
        "      res = decoder.predict(np.array([xi,xj]).reshape(-1,2))[0]\n",
        "      ax[i,j].imshow(res.reshape(28,28))\n",
        "      ax[i,j].axis('off')\n",
        "  plt.show()\n",
        "  \n",
        "plotsample(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WiKTnnIQUgdP"
      },
      "source": [
        "> **Task**: In our sample, we have trained fully-connected VAE. Now take the CNN from traditional auto-encoder above and create CNN-based VAE."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DhJLlh30Brpb"
      },
      "source": [
        "## Additional Materials\n",
        "\n",
        "* [Blog post on NeuroHive](https://neurohive.io/ru/osnovy-data-science/variacionnyj-avtojenkoder-vae/)\n",
        "* [Variational Autoencoders Explained](https://kvfrans.com/variational-autoencoders-explained/)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Autoencoders.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
