{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u2U7TIRdKNGu"
      },
      "source": [
        "# [Autoencoders](https://arxiv.org/abs/2201.03898)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "it5nmu_-c5-E"
      },
      "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",
        "Let's create simplest autoencoder for MNIST!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:39.262403Z",
          "iopub.status.busy": "2022-04-08T00:15:39.262149Z",
          "iopub.status.idle": "2022-04-08T00:15:39.268143Z",
          "shell.execute_reply": "2022-04-08T00:15:39.267475Z",
          "shell.execute_reply.started": "2022-04-08T00:15:39.262376Z"
        },
        "id": "6n8fzzN-2T1h",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import matplotlib.pyplot as plt\n",
        "from torchvision import transforms\n",
        "from torch import nn\n",
        "from torch import optim\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "torch.manual_seed(42)\n",
        "np.random.seed(42)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlAVPG-Ild19"
      },
      "source": [
        "Define training parameters and check if the GPU is available:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:39.530041Z",
          "iopub.status.busy": "2022-04-08T00:15:39.529522Z",
          "iopub.status.idle": "2022-04-08T00:15:39.534247Z",
          "shell.execute_reply": "2022-04-08T00:15:39.533327Z",
          "shell.execute_reply.started": "2022-04-08T00:15:39.530006Z"
        },
        "id": "bjL-jOgi3gmG",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "train_size = 0.9\n",
        "lr = 1e-3\n",
        "eps = 1e-8\n",
        "batch_size = 256\n",
        "epochs = 30"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jvsOL_ckld1-"
      },
      "source": [
        "The following function will load the MNIST dataset and apply specified transforms to it. It will also split it into train/test datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:39.893727Z",
          "iopub.status.busy": "2022-04-08T00:15:39.893160Z",
          "iopub.status.idle": "2022-04-08T00:15:39.899407Z",
          "shell.execute_reply": "2022-04-08T00:15:39.898565Z",
          "shell.execute_reply.started": "2022-04-08T00:15:39.893688Z"
        },
        "id": "g2Eo43713Sxb",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def mnist(train_part, transform=None):\n",
        "    dataset = torchvision.datasets.MNIST('.', download=True, transform=transform)\n",
        "    train_part = int(train_part * len(dataset))\n",
        "    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_part, len(dataset) - train_part])\n",
        "    return train_dataset, test_dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h86jOpd9ld1_"
      },
      "source": [
        "Now let's load the dataset and define dataloaders for train and test:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class CustomDataset(torch.utils.data.Dataset):\n",
        "    def __init__(self, raw_dataset, image_transform=None):\n",
        "        self.raw_dataset = raw_dataset\n",
        "        self.transform = image_transform\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.raw_dataset)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        raw = self.raw_dataset[idx][0]\n",
        "        target = raw.clone()\n",
        "        if self.transform is not None:\n",
        "            image = self.transform(raw.clone())\n",
        "        else:\n",
        "            image = raw.clone()\n",
        "        return image, target"
      ],
      "metadata": {
        "id": "MwBPD-rxmnAY"
      },
      "execution_count": 101,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transform = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "raw_train_dataset, raw_test_dataset = mnist(train_size, transform)\n",
        "train_dataset, test_dataset = CustomDataset(raw_train_dataset), CustomDataset(raw_test_dataset)"
      ],
      "metadata": {
        "id": "NNxiWHDDlstZ"
      },
      "execution_count": 102,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 103,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:41.292476Z",
          "iopub.status.busy": "2022-04-08T00:15:41.292238Z",
          "iopub.status.idle": "2022-04-08T00:15:41.299801Z",
          "shell.execute_reply": "2022-04-08T00:15:41.298897Z",
          "shell.execute_reply.started": "2022-04-08T00:15:41.292449Z"
        },
        "id": "jAI3uK86_zHM",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "train_dataloader = torch.utils.data.DataLoader(train_dataset, drop_last=True, batch_size=batch_size, shuffle=True)\n",
        "test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "dataloaders = (train_dataloader, test_dataloader)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:41.692950Z",
          "iopub.status.busy": "2022-04-08T00:15:41.692376Z",
          "iopub.status.idle": "2022-04-08T00:15:41.700707Z",
          "shell.execute_reply": "2022-04-08T00:15:41.699723Z",
          "shell.execute_reply.started": "2022-04-08T00:15:41.692909Z"
        },
        "id": "LdyQz4092fRV",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def plotn(n, data):\n",
        "    fig, ax = plt.subplots(1, n)\n",
        "    for i, z in enumerate(data):\n",
        "        if i == n:\n",
        "            break\n",
        "        preprocess = z[0].reshape(1, 28, 28) if z[0].shape[1] == 28 else z[0].reshape(1, 14, 14) if z[0].shape[1] == 14 else z[0]\n",
        "        ax[i].imshow(preprocess[0])\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 107,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:42.470833Z",
          "iopub.status.busy": "2022-04-08T00:15:42.470589Z",
          "iopub.status.idle": "2022-04-08T00:15:42.855919Z",
          "shell.execute_reply": "2022-04-08T00:15:42.853605Z",
          "shell.execute_reply.started": "2022-04-08T00:15:42.470808Z"
        },
        "id": "NeWJoiFC4A6J",
        "outputId": "e03b6437-3f38-49a2-c7bc-599322c04577",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 17100 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "plotn(5, train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 108,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:42.857823Z",
          "iopub.status.busy": "2022-04-08T00:15:42.857536Z",
          "iopub.status.idle": "2022-04-08T00:15:42.866008Z",
          "shell.execute_reply": "2022-04-08T00:15:42.865357Z",
          "shell.execute_reply.started": "2022-04-08T00:15:42.857787Z"
        },
        "id": "NnI2YvOg4DbT",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 16, kernel_size=(3, 3), padding='same')\n",
        "        self.maxpool1 = nn.MaxPool2d(kernel_size=(2, 2))\n",
        "        self.conv2 = nn.Conv2d(16, 8, kernel_size=(3, 3), padding='same')\n",
        "        self.maxpool2 = nn.MaxPool2d(kernel_size=(2, 2))\n",
        "        self.conv3 = nn.Conv2d(8, 8, kernel_size=(3, 3), padding='same')\n",
        "        self.maxpool3 = nn.MaxPool2d(kernel_size=(2, 2), padding=(1, 1))\n",
        "        self.relu = nn.ReLU()\n",
        "\n",
        "    def forward(self, input):\n",
        "        hidden1 = self.maxpool1(self.relu(self.conv1(input)))\n",
        "        hidden2 = self.maxpool2(self.relu(self.conv2(hidden1)))\n",
        "        encoded = self.maxpool3(self.relu(self.conv3(hidden2)))\n",
        "        return encoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 109,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:43.169659Z",
          "iopub.status.busy": "2022-04-08T00:15:43.169239Z",
          "iopub.status.idle": "2022-04-08T00:15:43.178700Z",
          "shell.execute_reply": "2022-04-08T00:15:43.178010Z",
          "shell.execute_reply.started": "2022-04-08T00:15:43.169622Z"
        },
        "id": "mZGB4Vr47478",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.conv1 = nn.Conv2d(8, 8, kernel_size=(3, 3), padding='same')\n",
        "        self.upsample1 = nn.Upsample(scale_factor=(2, 2))\n",
        "        self.conv2 = nn.Conv2d(8, 8, kernel_size=(3, 3), padding='same')\n",
        "        self.upsample2 = nn.Upsample(scale_factor=(2, 2))\n",
        "        self.conv3 = nn.Conv2d(8, 16, kernel_size=(3, 3))\n",
        "        self.upsample3 = nn.Upsample(scale_factor=(2, 2))\n",
        "        self.conv4 = nn.Conv2d(16, 1, kernel_size=(3, 3), padding='same')\n",
        "        self.relu = nn.ReLU()\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "\n",
        "    def forward(self, input):\n",
        "        hidden1 = self.upsample1(self.relu(self.conv1(input)))\n",
        "        hidden2 = self.upsample2(self.relu(self.conv2(hidden1)))\n",
        "        hidden3 = self.upsample3(self.relu(self.conv3(hidden2)))\n",
        "        decoded = self.sigmoid(self.conv4(hidden3))\n",
        "        return decoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 110,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:43.593770Z",
          "iopub.status.busy": "2022-04-08T00:15:43.593178Z",
          "iopub.status.idle": "2022-04-08T00:15:43.600152Z",
          "shell.execute_reply": "2022-04-08T00:15:43.599205Z",
          "shell.execute_reply.started": "2022-04-08T00:15:43.593731Z"
        },
        "id": "SDGiAPhbBBLY",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class AutoEncoder(nn.Module):\n",
        "    def __init__(self, super_resolution=False):\n",
        "        super().__init__()\n",
        "        if not super_resolution:\n",
        "            self.encoder = Encoder()\n",
        "        else:\n",
        "            self.encoder = SuperResolutionEncoder()\n",
        "        self.decoder = Decoder()\n",
        "\n",
        "    def forward(self, input):\n",
        "        encoded = self.encoder(input)\n",
        "        decoded = self.decoder(encoded)\n",
        "        return decoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 111,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:44.069329Z",
          "iopub.status.busy": "2022-04-08T00:15:44.068875Z",
          "iopub.status.idle": "2022-04-08T00:15:44.079486Z",
          "shell.execute_reply": "2022-04-08T00:15:44.078748Z",
          "shell.execute_reply.started": "2022-04-08T00:15:44.069289Z"
        },
        "id": "nZyG_mNu_Pnc",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "model = AutoEncoder().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, eps=eps)\n",
        "loss_fn = nn.BCELoss()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 112,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T01:25:58.828126Z",
          "iopub.status.busy": "2022-04-08T01:25:58.827634Z",
          "iopub.status.idle": "2022-04-08T01:25:58.844248Z",
          "shell.execute_reply": "2022-04-08T01:25:58.843270Z",
          "shell.execute_reply.started": "2022-04-08T01:25:58.828092Z"
        },
        "id": "iiIy87v2_rUr",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "def train(dataloaders, model, loss_fn, optimizer, epochs, device):\n",
        "    tqdm_iter = tqdm(range(epochs))\n",
        "    train_dataloader, test_dataloader = dataloaders[0], dataloaders[1]\n",
        "\n",
        "    for epoch in tqdm_iter:\n",
        "        model.train()\n",
        "        train_loss = 0.0\n",
        "        test_loss = 0.0\n",
        "\n",
        "        for batch in train_dataloader:\n",
        "            imgs, labels = batch\n",
        "            imgs = imgs.to(device)\n",
        "            labels = labels.to(device)\n",
        "\n",
        "            preds = model(imgs)\n",
        "            loss = loss_fn(preds, labels)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            train_loss += loss.item()\n",
        "\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            for batch in test_dataloader:\n",
        "                imgs, labels = batch\n",
        "                imgs = imgs.to(device)\n",
        "                labels = labels.to(device)\n",
        "\n",
        "                preds = model(imgs)\n",
        "                loss = loss_fn(preds, labels)\n",
        "\n",
        "                test_loss += loss.item()\n",
        "\n",
        "        train_loss /= len(train_dataloader)\n",
        "        test_loss /= len(test_dataloader)\n",
        "\n",
        "        tqdm_dct = {'train loss:': train_loss, 'test loss:': test_loss}\n",
        "        tqdm_iter.set_postfix(tqdm_dct, refresh=True)\n",
        "        tqdm_iter.refresh()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 113,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:15:44.969616Z",
          "iopub.status.busy": "2022-04-08T00:15:44.969345Z",
          "iopub.status.idle": "2022-04-08T00:22:34.533469Z",
          "shell.execute_reply": "2022-04-08T00:22:34.532667Z",
          "shell.execute_reply.started": "2022-04-08T00:15:44.969587Z"
        },
        "id": "PMqO8eOxCemz",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e419b7c6-1966-4d8e-ec31-db1e888d4884"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 30/30 [14:48<00:00, 29.61s/it, train loss:=0.104, test loss:=0.104]\n"
          ]
        }
      ],
      "source": [
        "train(dataloaders, model, loss_fn, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 114,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:22:34.535491Z",
          "iopub.status.busy": "2022-04-08T00:22:34.535154Z",
          "iopub.status.idle": "2022-04-08T00:22:35.701465Z",
          "shell.execute_reply": "2022-04-08T00:22:35.700797Z",
          "shell.execute_reply.started": "2022-04-08T00:22:34.535452Z"
        },
        "id": "kR3n0EnjOts0",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "outputId": "403019d2-b278-4804-c386-157bf908a50a"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 17516 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 21924 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "plots = 5\n",
        "for i, data in enumerate(test_dataset):\n",
        "    if i == plots:\n",
        "        break\n",
        "    predictions.append(model(data[0].to(device).unsqueeze(0)).detach().cpu())\n",
        "plotn(plots, test_dataset)\n",
        "plotn(plots, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1JUmf9i1dg2i"
      },
      "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": "tjdFt03rULX-"
      },
      "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",
      "source": [
        "def noisy_transform(x):\n",
        "    return np.clip(x + torch.FloatTensor(np.random.normal(loc=0.5, scale=0.3, size=x.shape)), 0.0, 1.0)"
      ],
      "metadata": {
        "id": "wQ66Jr6QqOSl"
      },
      "execution_count": 131,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transform = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "raw_train_dataset, raw_test_dataset = mnist(train_size, transform)\n",
        "train_dataset, test_dataset = CustomDataset(raw_train_dataset, image_transform=noisy_transform), CustomDataset(raw_test_dataset, image_transform=noisy_transform)"
      ],
      "metadata": {
        "id": "YkBqWagdqB3w"
      },
      "execution_count": 132,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataloader = torch.utils.data.DataLoader(train_dataset, drop_last=True, batch_size=batch_size, shuffle=True)\n",
        "test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "dataloaders = (train_dataloader, test_dataloader)"
      ],
      "metadata": {
        "id": "gmimJVV0txQz"
      },
      "execution_count": 133,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 134,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T01:26:06.259333Z",
          "iopub.status.busy": "2022-04-08T01:26:06.259061Z",
          "iopub.status.idle": "2022-04-08T01:26:06.799264Z",
          "shell.execute_reply": "2022-04-08T01:26:06.798609Z",
          "shell.execute_reply.started": "2022-04-08T01:26:06.259303Z"
        },
        "id": "1Yj9ZRDmUPxX",
        "outputId": "5267bef1-e30a-47d3-f62f-b4594a3e26b7",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 62316 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "plotn(5, train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 135,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T01:55:45.901599Z",
          "iopub.status.busy": "2022-04-08T01:55:45.901279Z",
          "iopub.status.idle": "2022-04-08T01:55:45.910993Z",
          "shell.execute_reply": "2022-04-08T01:55:45.909858Z",
          "shell.execute_reply.started": "2022-04-08T01:55:45.901533Z"
        },
        "id": "qxo8NDLLUvut",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "model = AutoEncoder().to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, eps=eps)\n",
        "loss_fn = nn.BCELoss()\n",
        "epochs = 100"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 136,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T01:55:47.956516Z",
          "iopub.status.busy": "2022-04-08T01:55:47.956250Z",
          "iopub.status.idle": "2022-04-08T02:18:17.428958Z",
          "shell.execute_reply": "2022-04-08T02:18:17.428239Z",
          "shell.execute_reply.started": "2022-04-08T01:55:47.956489Z"
        },
        "id": "JAYzgfoTUBHM",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a7d2d89e-f35c-45de-b2c3-f0ca8e81031a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [1:19:55<00:00, 47.96s/it, train loss:=0.112, test loss:=0.113]\n"
          ]
        }
      ],
      "source": [
        "train(dataloaders, model, loss_fn, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 137,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T02:18:17.430982Z",
          "iopub.status.busy": "2022-04-08T02:18:17.430570Z",
          "iopub.status.idle": "2022-04-08T02:18:18.235054Z",
          "shell.execute_reply": "2022-04-08T02:18:18.234364Z",
          "shell.execute_reply.started": "2022-04-08T02:18:17.430943Z"
        },
        "id": "IaPfyJ0SV7XY",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "outputId": "87afb9b8-fbde-4223-dfcc-72dd0e2b6b56"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 62240 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 19700 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "plots = 5\n",
        "for i, data in enumerate(test_dataset):\n",
        "    if i == plots:\n",
        "        break\n",
        "    predictions.append(model(data[0].to(device).unsqueeze(0)).detach().cpu())\n",
        "plotn(plots, test_dataset)\n",
        "plotn(plots, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MyvmZEHzdoJT"
      },
      "source": [
        "> **Exercise:** See how denoiser trained on MNIST digits works for different images. As an example, you can take [Fashion MNIST](https://pytorch.org/vision/stable/generated/torchvision.datasets.FashionMNIST.html#torchvision.datasets.FashionMNIST) 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": "yxFKm_OxdqfO"
      },
      "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",
        "For that let's downscale image to 14x14 at train."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "transform = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "raw_train_dataset, raw_test_dataset = mnist(train_size, transform)\n",
        "\n",
        "super_res_transform = transforms.Resize((14, 14))\n",
        "train_dataset, test_dataset = CustomDataset(raw_train_dataset, image_transform=super_res_transform), CustomDataset(raw_test_dataset, image_transform=super_res_transform)"
      ],
      "metadata": {
        "id": "9tmwSp-frz8X"
      },
      "execution_count": 138,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataloader = torch.utils.data.DataLoader(train_dataset, drop_last=True, batch_size=batch_size, shuffle=True)\n",
        "test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "dataloaders = (train_dataloader, test_dataloader)"
      ],
      "metadata": {
        "id": "p5NrDP5Jt0M7"
      },
      "execution_count": 139,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 140,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:31:45.310634Z",
          "iopub.status.busy": "2022-04-08T00:31:45.310384Z",
          "iopub.status.idle": "2022-04-08T00:31:45.688313Z",
          "shell.execute_reply": "2022-04-08T00:31:45.687614Z",
          "shell.execute_reply.started": "2022-04-08T00:31:45.310597Z"
        },
        "id": "trida5guu9js",
        "outputId": "62ba84e5-ff3d-4f56-a576-21e21e322e68",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 108
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 4844 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "plotn(5, train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 141,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:31:45.690022Z",
          "iopub.status.busy": "2022-04-08T00:31:45.689620Z",
          "iopub.status.idle": "2022-04-08T00:31:45.698351Z",
          "shell.execute_reply": "2022-04-08T00:31:45.697586Z",
          "shell.execute_reply.started": "2022-04-08T00:31:45.689984Z"
        },
        "id": "9vC57e-rei4p",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class SuperResolutionEncoder(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 16, kernel_size=(3, 3), padding='same')\n",
        "        self.maxpool1 = nn.MaxPool2d(kernel_size=(2, 2))\n",
        "        self.conv2 = nn.Conv2d(16, 8, kernel_size=(3, 3), padding='same')\n",
        "        self.maxpool2 = nn.MaxPool2d(kernel_size=(2, 2), padding=(1, 1))\n",
        "        self.relu = nn.ReLU()\n",
        "\n",
        "    def forward(self, input):\n",
        "        hidden1 = self.maxpool1(self.relu(self.conv1(input)))\n",
        "        encoded = self.maxpool2(self.relu(self.conv2(hidden1)))\n",
        "        return encoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 142,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:31:45.699896Z",
          "iopub.status.busy": "2022-04-08T00:31:45.699576Z",
          "iopub.status.idle": "2022-04-08T00:31:45.714684Z",
          "shell.execute_reply": "2022-04-08T00:31:45.713859Z",
          "shell.execute_reply.started": "2022-04-08T00:31:45.699857Z"
        },
        "id": "d78J288qe5qJ",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "model = AutoEncoder(super_resolution=True).to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, eps=eps)\n",
        "loss_fn = nn.BCELoss()\n",
        "epochs = 30"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 143,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:31:45.718432Z",
          "iopub.status.busy": "2022-04-08T00:31:45.718047Z",
          "iopub.status.idle": "2022-04-08T00:38:29.683824Z",
          "shell.execute_reply": "2022-04-08T00:38:29.683115Z",
          "shell.execute_reply.started": "2022-04-08T00:31:45.718402Z"
        },
        "id": "CJ_zcN5Je6I-",
        "outputId": "943bcc64-f818-4c71-f1e1-ce70883b29a2",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 30/30 [15:12<00:00, 30.43s/it, train loss:=0.104, test loss:=0.104]\n"
          ]
        }
      ],
      "source": [
        "train(dataloaders, model, loss_fn, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 144,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:29.687844Z",
          "iopub.status.busy": "2022-04-08T00:38:29.687353Z",
          "iopub.status.idle": "2022-04-08T00:38:30.585026Z",
          "shell.execute_reply": "2022-04-08T00:38:30.584352Z",
          "shell.execute_reply.started": "2022-04-08T00:38:29.687803Z"
        },
        "id": "YsVfcCKKfjv1",
        "outputId": "b31eeeaf-97fd-452d-ace6-6b6bd24156d4",
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 200
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 4844 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ],
            "image/png": "[encoded data, 20476 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "plots = 5\n",
        "for i, data in enumerate(test_dataset):\n",
        "    if i == plots:\n",
        "        break\n",
        "    predictions.append(model(data[0].to(device).unsqueeze(0)).detach().cpu())\n",
        "plotn(plots, test_dataset)\n",
        "plotn(plots, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-aZqFJthfu9-"
      },
      "source": [
        "> **Exercise**: Try to train super-resolution network on [CIFAR-10](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html) for 2x and 4x upscaling. Use noise as input to 4x upscaling model and observe the result."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A3DOJU1-gTJV"
      },
      "source": [
        "# [Variational Auto-Encoders (VAE)](https://arxiv.org/abs/1906.02691)\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}})$, where $\\mathrm{z\\_mean}, \\mathrm{z\\_log} \\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` (instead of predicting the standard deviation itself, we predict it's logarithm)\n",
        " * We sample a vector `sample(z_val in code)` 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%\">\n",
        "\n",
        " > Image from [this blog post](https://ijdykeman.github.io/ml/2016/12/21/cvae.html) by Isaak Dykeman"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.586882Z",
          "iopub.status.busy": "2022-04-08T00:38:30.586394Z",
          "iopub.status.idle": "2022-04-08T00:38:30.595011Z",
          "shell.execute_reply": "2022-04-08T00:38:30.594391Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.586844Z"
        },
        "id": "aT_GEWeU409I",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class VAEEncoder(nn.Module):\n",
        "    def __init__(self, device):\n",
        "        super().__init__()\n",
        "        self.intermediate_dim = 512\n",
        "        self.latent_dim = 2\n",
        "        self.linear = nn.Linear(784, self.intermediate_dim)\n",
        "        self.z_mean = nn.Linear(self.intermediate_dim, self.latent_dim)\n",
        "        self.z_log = nn.Linear(self.intermediate_dim, self.latent_dim)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.device = device\n",
        "\n",
        "    def forward(self, input):\n",
        "        bs = input.shape[0]\n",
        "\n",
        "        hidden = self.relu(self.linear(input))\n",
        "        z_mean = self.z_mean(hidden)\n",
        "        z_log = self.z_log(hidden)\n",
        "\n",
        "        eps = torch.FloatTensor(np.random.normal(size=(bs, self.latent_dim))).to(device)\n",
        "        z_val = z_mean + torch.exp(z_log) * eps\n",
        "        return z_mean, z_log, z_val"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.596177Z",
          "iopub.status.busy": "2022-04-08T00:38:30.595932Z",
          "iopub.status.idle": "2022-04-08T00:38:30.606805Z",
          "shell.execute_reply": "2022-04-08T00:38:30.606054Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.596144Z"
        },
        "id": "XWi4oCcq409p",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class VAEDecoder(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.intermediate_dim = 512\n",
        "        self.latent_dim = 2\n",
        "        self.linear = nn.Linear(self.latent_dim, self.intermediate_dim)\n",
        "        self.output = nn.Linear(self.intermediate_dim, 784)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "\n",
        "    def forward(self, input):\n",
        "        hidden = self.relu(self.linear(input))\n",
        "        decoded = self.sigmoid(self.output(hidden))\n",
        "        return decoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.608150Z",
          "iopub.status.busy": "2022-04-08T00:38:30.607836Z",
          "iopub.status.idle": "2022-04-08T00:38:30.618464Z",
          "shell.execute_reply": "2022-04-08T00:38:30.617740Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.608114Z"
        },
        "id": "WukDYQ9f409p",
        "trusted": true
      },
      "outputs": [],
      "source": [
        "class VAEAutoEncoder(nn.Module):\n",
        "    def __init__(self, device):\n",
        "        super().__init__()\n",
        "        self.encoder = VAEEncoder(device)\n",
        "        self.decoder = VAEDecoder()\n",
        "        self.z_vals = None\n",
        "\n",
        "    def forward(self, input):\n",
        "        bs, c, h, w = input.shape[0], input.shape[1], input.shape[2], input.shape[3]\n",
        "        input = input.view(bs, -1)\n",
        "        encoded = self.encoder(input)\n",
        "        self.z_vals = encoded\n",
        "        decoded = self.decoder(encoded[2])\n",
        "        return decoded\n",
        "    \n",
        "    def get_zvals(self):\n",
        "        return self.z_vals"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EiNMancEld2J"
      },
      "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": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.620121Z",
          "iopub.status.busy": "2022-04-08T00:38:30.619868Z",
          "iopub.status.idle": "2022-04-08T00:38:30.627984Z",
          "shell.execute_reply": "2022-04-08T00:38:30.627265Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.620089Z"
        },
        "trusted": true,
        "id": "ZBojFRTPld2J"
      },
      "outputs": [],
      "source": [
        "def vae_loss(preds, targets, z_vals):\n",
        "    mse = nn.MSELoss()\n",
        "    reconstruction_loss = mse(preds, targets.view(targets.shape[0], -1)) * 784.0\n",
        "    temp = 1.0 + z_vals[1] - torch.square(z_vals[0]) - torch.exp(z_vals[1])\n",
        "    kl_loss = -0.5 * torch.sum(temp, axis=-1)\n",
        "    return torch.mean(reconstruction_loss + kl_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.629709Z",
          "iopub.status.busy": "2022-04-08T00:38:30.629192Z",
          "iopub.status.idle": "2022-04-08T00:38:30.648407Z",
          "shell.execute_reply": "2022-04-08T00:38:30.647801Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.629671Z"
        },
        "trusted": true,
        "id": "0-iAt8R-ld2K"
      },
      "outputs": [],
      "source": [
        "model = VAEAutoEncoder(device).to(device)\n",
        "optimizer = optim.RMSprop(model.parameters(), lr=lr, eps=eps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.649871Z",
          "iopub.status.busy": "2022-04-08T00:38:30.649560Z",
          "iopub.status.idle": "2022-04-08T00:38:30.659437Z",
          "shell.execute_reply": "2022-04-08T00:38:30.658769Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.649835Z"
        },
        "trusted": true,
        "id": "3z4fkBdKld2K"
      },
      "outputs": [],
      "source": [
        "def train_vae(dataloaders, model, optimizer, epochs, device):\n",
        "    tqdm_iter = tqdm(range(epochs))\n",
        "    train_dataloader, test_dataloader = dataloaders[0], dataloaders[1]\n",
        "\n",
        "    for epoch in tqdm_iter:\n",
        "        model.train()\n",
        "        train_loss = 0.0\n",
        "        test_loss = 0.0\n",
        "\n",
        "        for batch in train_dataloader:\n",
        "            imgs, labels = batch\n",
        "            imgs = imgs.to(device)\n",
        "            labels = labels.to(device)\n",
        "\n",
        "            preds = model(imgs)\n",
        "            z_vals = model.get_zvals()\n",
        "            loss = vae_loss(preds, imgs, z_vals)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            train_loss += loss.item()\n",
        "\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            for batch in test_dataloader:\n",
        "                imgs, labels = batch\n",
        "                imgs = imgs.to(device)\n",
        "                labels = labels.to(device)\n",
        "\n",
        "                preds = model(imgs)\n",
        "                z_vals = model.get_zvals()\n",
        "                loss = vae_loss(preds, imgs, z_vals)\n",
        "\n",
        "                test_loss += loss.item()\n",
        "\n",
        "        train_loss /= len(train_dataloader)\n",
        "        test_loss /= len(test_dataloader)\n",
        "\n",
        "        tqdm_dct = {'train loss:': train_loss, 'test loss:': test_loss}\n",
        "        tqdm_iter.set_postfix(tqdm_dct, refresh=True)\n",
        "        tqdm_iter.refresh()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:38:30.661256Z",
          "iopub.status.busy": "2022-04-08T00:38:30.660812Z",
          "iopub.status.idle": "2022-04-08T00:43:25.674280Z",
          "shell.execute_reply": "2022-04-08T00:43:25.673608Z",
          "shell.execute_reply.started": "2022-04-08T00:38:30.661216Z"
        },
        "trusted": true,
        "id": "XaA3HRDlld2K",
        "outputId": "e91b9a6d-e15c-4b89-fc0e-6f126c42252c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 30/30 [04:54<00:00,  9.83s/it, train loss:=35.1, test loss:=35.6]\n"
          ]
        }
      ],
      "source": [
        "train_vae(dataloaders, model, optimizer, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:25.676045Z",
          "iopub.status.busy": "2022-04-08T00:43:25.675634Z",
          "iopub.status.idle": "2022-04-08T00:43:26.429255Z",
          "shell.execute_reply": "2022-04-08T00:43:26.428437Z",
          "shell.execute_reply.started": "2022-04-08T00:43:25.676007Z"
        },
        "trusted": true,
        "id": "xCTlw0PJld2K",
        "outputId": "db0f2905-78b2-491a-e4bc-c4c2ef2103a0"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 17480 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 24520 characters]",
            "text/plain": [
              "<Figure size 432x288 with 5 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "model.eval()\n",
        "predictions = []\n",
        "plots = 5\n",
        "for i, data in enumerate(test_dataset):\n",
        "    if i == plots:\n",
        "        break\n",
        "    predictions.append(model(data[0].to(device).unsqueeze(0)).view(1, 28, 28).detach().cpu())\n",
        "plotn(plots, test_dataset)\n",
        "plotn(plots, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CgGF6hjZld2L"
      },
      "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": "KydrDl8Ild2L"
      },
      "source": [
        "# [Adversarial Auto-Encoders (AAE)](https://arxiv.org/abs/1511.05644)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFspsM54ld2L"
      },
      "source": [
        "Adversarial Auto-Encoders is a **combination** of Generative Adversarial Networks and Variational Auto-Encoders. \n",
        "\n",
        "Encoder will be the generator, discriminator will learn to distinguish the real images encoder output from generated ones. Encoder output is a distribution, from this output decoder will try decode image.\n",
        "\n",
        "In this approach we have **three loss functions**: generator loss, discriminator loss from GAN's and reconstruction loss from VAE.\n",
        "\n",
        " <img src=\"images/aae.png\" width=\"50%\">\n",
        "\n",
        " > Image from [this blog post](https://blog.paperspace.com/adversarial-autoencoders-with-pytorch/) by Felipe Ducau"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.430881Z",
          "iopub.status.busy": "2022-04-08T00:43:26.430435Z",
          "iopub.status.idle": "2022-04-08T00:43:26.437703Z",
          "shell.execute_reply": "2022-04-08T00:43:26.437030Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.430837Z"
        },
        "trusted": true,
        "id": "bCktg2nEld2L"
      },
      "outputs": [],
      "source": [
        "class AAEEncoder(nn.Module):\n",
        "    def __init__(self, input_dim, inter_dim, latent_dim):\n",
        "        super().__init__()\n",
        "        self.linear1 = nn.Linear(input_dim, inter_dim)\n",
        "        self.linear2 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear3 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear4 = nn.Linear(inter_dim, latent_dim)\n",
        "        self.relu = nn.ReLU()\n",
        "        \n",
        "    def forward(self, input):\n",
        "        hidden1 = self.relu(self.linear1(input))\n",
        "        hidden2 = self.relu(self.linear2(hidden1))\n",
        "        hidden3 = self.relu(self.linear3(hidden2))\n",
        "        encoded = self.linear4(hidden3)\n",
        "        return encoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.439383Z",
          "iopub.status.busy": "2022-04-08T00:43:26.438973Z",
          "iopub.status.idle": "2022-04-08T00:43:26.450914Z",
          "shell.execute_reply": "2022-04-08T00:43:26.450120Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.439339Z"
        },
        "trusted": true,
        "id": "ygqUQvqhld2L"
      },
      "outputs": [],
      "source": [
        "class AAEDecoder(nn.Module):\n",
        "    def __init__(self, latent_dim, inter_dim, output_dim):\n",
        "        super().__init__()\n",
        "        self.linear1 = nn.Linear(latent_dim, inter_dim)\n",
        "        self.linear2 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear3 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear4 = nn.Linear(inter_dim, output_dim)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "        \n",
        "    def forward(self, input):\n",
        "        hidden1 = self.relu(self.linear1(input))\n",
        "        hidden2 = self.relu(self.linear2(hidden1))\n",
        "        hidden3 = self.relu(self.linear3(hidden2))\n",
        "        decoded = self.sigmoid(self.linear4(hidden3))\n",
        "        return decoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.453838Z",
          "iopub.status.busy": "2022-04-08T00:43:26.453577Z",
          "iopub.status.idle": "2022-04-08T00:43:26.462584Z",
          "shell.execute_reply": "2022-04-08T00:43:26.461884Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.453814Z"
        },
        "trusted": true,
        "id": "cZhHZeCsld2L"
      },
      "outputs": [],
      "source": [
        "class AAEDiscriminator(nn.Module):\n",
        "    def __init__(self, latent_dim, inter_dim):\n",
        "        super().__init__()\n",
        "        self.latent_dim = latent_dim\n",
        "        self.inter_dim = inter_dim\n",
        "        self.linear1 = nn.Linear(latent_dim, inter_dim)\n",
        "        self.linear2 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear3 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear4 = nn.Linear(inter_dim, inter_dim)\n",
        "        self.linear5 = nn.Linear(inter_dim, 1)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "        \n",
        "    def forward(self, input):\n",
        "        hidden1 = self.relu(self.linear1(input))\n",
        "        hidden2 = self.relu(self.linear2(hidden1))\n",
        "        hidden3 = self.relu(self.linear3(hidden2))\n",
        "        hidden4 = self.relu(self.linear4(hidden3))\n",
        "        decoded = self.sigmoid(self.linear4(hidden4))\n",
        "        return decoded\n",
        "    \n",
        "    def get_dims(self):\n",
        "        return self.latent_dim, self.inter_dim\n",
        "        "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.464666Z",
          "iopub.status.busy": "2022-04-08T00:43:26.464215Z",
          "iopub.status.idle": "2022-04-08T00:43:26.474821Z",
          "shell.execute_reply": "2022-04-08T00:43:26.474117Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.464628Z"
        },
        "trusted": true,
        "id": "xb3KXNK8ld2M"
      },
      "outputs": [],
      "source": [
        "input_dims = 784\n",
        "inter_dims = 1000\n",
        "latent_dims = 150"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.477879Z",
          "iopub.status.busy": "2022-04-08T00:43:26.477199Z",
          "iopub.status.idle": "2022-04-08T00:43:26.537720Z",
          "shell.execute_reply": "2022-04-08T00:43:26.537052Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.477797Z"
        },
        "trusted": true,
        "id": "vdZInPOQld2M"
      },
      "outputs": [],
      "source": [
        "aae_encoder = AAEEncoder(input_dims, inter_dims, latent_dims).to(device)\n",
        "aae_decoder = AAEDecoder(latent_dims, inter_dims, input_dims).to(device)\n",
        "aae_discriminator = AAEDiscriminator(latent_dims, int(inter_dims / 2)).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.539592Z",
          "iopub.status.busy": "2022-04-08T00:43:26.539109Z",
          "iopub.status.idle": "2022-04-08T00:43:26.543949Z",
          "shell.execute_reply": "2022-04-08T00:43:26.543156Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.539555Z"
        },
        "trusted": true,
        "id": "qYhyBPDqld2M"
      },
      "outputs": [],
      "source": [
        "lr = 1e-4\n",
        "regularization_lr = 5e-5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.545858Z",
          "iopub.status.busy": "2022-04-08T00:43:26.545404Z",
          "iopub.status.idle": "2022-04-08T00:43:26.553377Z",
          "shell.execute_reply": "2022-04-08T00:43:26.552700Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.545825Z"
        },
        "trusted": true,
        "id": "Di0y1zVOld2M"
      },
      "outputs": [],
      "source": [
        "optim_encoder = optim.Adam(aae_encoder.parameters(), lr=lr)\n",
        "optim_encoder_regularization = optim.Adam(aae_encoder.parameters(), lr=regularization_lr)\n",
        "optim_decoder = optim.Adam(aae_decoder.parameters(), lr=lr)\n",
        "optim_discriminator = optim.Adam(aae_discriminator.parameters(), lr=regularization_lr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.555945Z",
          "iopub.status.busy": "2022-04-08T00:43:26.555659Z",
          "iopub.status.idle": "2022-04-08T00:43:26.576630Z",
          "shell.execute_reply": "2022-04-08T00:43:26.575929Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.555908Z"
        },
        "trusted": true,
        "id": "_1ClSDv1ld2M"
      },
      "outputs": [],
      "source": [
        "def train_aae(dataloaders, models, optimizers, epochs, device):\n",
        "    tqdm_iter = tqdm(range(epochs))\n",
        "    train_dataloader, test_dataloader = dataloaders[0], dataloaders[1]\n",
        "    \n",
        "    enc, dec, disc = models[0], models[1], models[2]\n",
        "    optim_enc, optim_enc_reg, optim_dec, optim_disc = optimizers[0], optimizers[1], optimizers[2], optimizers[3]\n",
        "    \n",
        "    eps = 1e-9\n",
        "\n",
        "    for epoch in tqdm_iter:\n",
        "        enc.train()\n",
        "        dec.train()\n",
        "        disc.train()\n",
        "\n",
        "        train_reconst_loss = 0.0\n",
        "        train_disc_loss = 0.0\n",
        "        train_enc_loss = 0.0\n",
        "        \n",
        "        test_reconst_loss = 0.0\n",
        "        test_disc_loss = 0.0\n",
        "        test_enc_loss = 0.0\n",
        "\n",
        "        for batch in train_dataloader:\n",
        "            imgs, labels = batch\n",
        "            imgs = imgs.view(imgs.shape[0], -1).to(device)\n",
        "            labels = labels.to(device)\n",
        "            \n",
        "            enc.zero_grad()\n",
        "            dec.zero_grad()\n",
        "            disc.zero_grad()\n",
        "             \n",
        "            encoded = enc(imgs)\n",
        "            decoded = dec(encoded)\n",
        "            \n",
        "            reconstruction_loss = F.binary_cross_entropy(decoded, imgs)\n",
        "            reconstruction_loss.backward()\n",
        "            \n",
        "            optim_enc.step()\n",
        "            optim_dec.step()\n",
        "            enc.eval()\n",
        "\n",
        "            latent_dim, disc_inter_dim = disc.get_dims()\n",
        "            real = torch.randn(imgs.shape[0], latent_dim).to(device)\n",
        "            \n",
        "            disc_real = disc(real)\n",
        "            disc_fake = disc(enc(imgs))\n",
        "            \n",
        "            disc_loss = -torch.mean(torch.log(disc_real + eps) + torch.log(1.0 - disc_fake + eps))\n",
        "            disc_loss.backward()\n",
        "            \n",
        "            optim_dec.step()\n",
        "            enc.train()\n",
        "            \n",
        "            disc_fake = disc(enc(imgs))\n",
        "            enc_loss = -torch.mean(torch.log(disc_fake + eps))\n",
        "            enc_loss.backward()\n",
        "            \n",
        "            optim_enc_reg.step()\n",
        "\n",
        "            train_reconst_loss += reconstruction_loss.item()\n",
        "            train_disc_loss += disc_loss.item()\n",
        "            train_enc_loss += enc_loss.item()\n",
        "\n",
        "        enc.eval()\n",
        "        dec.eval()\n",
        "        disc.eval()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for batch in test_dataloader:\n",
        "                imgs, labels = batch\n",
        "                imgs = imgs.view(imgs.shape[0], -1).to(device)\n",
        "                labels = labels.to(device)\n",
        "\n",
        "                encoded = enc(imgs)\n",
        "                decoded = dec(encoded)\n",
        "\n",
        "                reconstruction_loss = F.binary_cross_entropy(decoded, imgs)\n",
        "\n",
        "                latent_dim, disc_inter_dim = disc.get_dims()\n",
        "                real = torch.randn(imgs.shape[0], latent_dim).to(device)\n",
        "\n",
        "                disc_real = disc(real)\n",
        "                disc_fake = disc(enc(imgs))\n",
        "                disc_loss = -torch.mean(torch.log(disc_real + eps) + torch.log(1.0 - disc_fake + eps))\n",
        "\n",
        "                disc_fake = disc(enc(imgs))\n",
        "                enc_loss = -torch.mean(torch.log(disc_fake + eps))\n",
        "\n",
        "                test_reconst_loss += reconstruction_loss.item()\n",
        "                test_disc_loss += disc_loss.item()\n",
        "                test_enc_loss += enc_loss.item()\n",
        "\n",
        "        train_reconst_loss /= len(train_dataloader)\n",
        "        train_disc_loss /= len(train_dataloader)\n",
        "        train_enc_loss /= len(train_dataloader)\n",
        "        \n",
        "        test_reconst_loss /= len(test_dataloader)\n",
        "        test_disc_loss /= len(test_dataloader)\n",
        "        test_enc_loss /= len(test_dataloader)\n",
        "\n",
        "        tqdm_dct = {'train reconst loss:': train_reconst_loss, 'train disc loss:': train_disc_loss, 'train enc loss': train_enc_loss, \\\n",
        "                        'test reconst loss:': test_reconst_loss, 'test disc loss:': test_disc_loss, 'test enc loss': test_enc_loss}\n",
        "        tqdm_iter.set_postfix(tqdm_dct, refresh=True)\n",
        "        tqdm_iter.refresh()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.578419Z",
          "iopub.status.busy": "2022-04-08T00:43:26.577933Z",
          "iopub.status.idle": "2022-04-08T00:43:26.589197Z",
          "shell.execute_reply": "2022-04-08T00:43:26.588532Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.578373Z"
        },
        "trusted": true,
        "id": "0Z7Qt2R_ld2M"
      },
      "outputs": [],
      "source": [
        "models = (aae_encoder, aae_decoder, aae_discriminator)\n",
        "optimizers = (optim_encoder, optim_encoder_regularization, optim_decoder, optim_discriminator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:43:26.592138Z",
          "iopub.status.busy": "2022-04-08T00:43:26.590197Z",
          "iopub.status.idle": "2022-04-08T00:52:49.002321Z",
          "shell.execute_reply": "2022-04-08T00:52:49.001592Z",
          "shell.execute_reply.started": "2022-04-08T00:43:26.592099Z"
        },
        "trusted": true,
        "id": "MeWnsKPSld2N",
        "outputId": "20be03ca-d9b1-44c6-c10e-f069b0d9b441"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 30/30 [09:22<00:00, 18.75s/it, train reconst loss:=0.0919, train disc loss:=1.39, train enc loss=0.692, test reconst loss:=0.0945, test disc loss:=1.39, test enc loss=0.692]\n"
          ]
        }
      ],
      "source": [
        "train_aae(dataloaders, models, optimizers, epochs, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "execution": {
          "iopub.execute_input": "2022-04-08T00:52:49.004318Z",
          "iopub.status.busy": "2022-04-08T00:52:49.003842Z",
          "iopub.status.idle": "2022-04-08T00:52:50.737055Z",
          "shell.execute_reply": "2022-04-08T00:52:50.736379Z",
          "shell.execute_reply.started": "2022-04-08T00:52:49.004280Z"
        },
        "trusted": true,
        "id": "Ddt97NZYld2N",
        "outputId": "0d104380-98a2-4249-a2a3-875143825db2"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 12944 characters]",
            "text/plain": [
              "<Figure size 432x288 with 10 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        },
        {
          "data": {
            "image/png": "[encoded data, 14884 characters]",
            "text/plain": [
              "<Figure size 432x288 with 10 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "aae_encoder.eval()\n",
        "aae_decoder.eval()\n",
        "predictions = []\n",
        "plots = 10\n",
        "for i, data in enumerate(test_dataset):\n",
        "    if i == plots:\n",
        "        break\n",
        "    pred = aae_decoder(aae_encoder(data[0].to(device).unsqueeze(0).view(1, 784)))\n",
        "    predictions.append(pred.view(1, 28, 28).detach().cpu())\n",
        "plotn(plots, test_dataset)\n",
        "plotn(plots, predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owD_X5aPld2N"
      },
      "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": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.12"
    },
    "colab": {
      "name": "AutoEncodersPyTorch.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "yxFKm_OxdqfO",
        "A3DOJU1-gTJV",
        "KydrDl8Ild2L",
        "owD_X5aPld2N"
      ]
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}