homeworks/hw2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CyfGcNY4pcXd"
   },
   "source": [
    "# Homework2: Autoregressive models & VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EIBqEphlrEGd",
    "outputId": "28bb7352-1420-4af5-d5ac-ffcf0b02f740"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "\n",
    "print('cuda is available:', USE_CUDA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0wq164ymr-68",
    "outputId": "0a4e5012-c569-4dae-9cfa-850b06d9e722"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gf3_Sb-Pq_L8"
   },
   "outputs": [],
   "source": [
    "def load_pickle(path, flatten=True):\n",
    "    with open(path, 'rb') as f:\n",
    "        data = pickle.load(f)\n",
    "    train_data = data['train'].astype('float32')[:, :, :, [0]] > 128\n",
    "    test_data = data['test'].astype('float32')[:, :, :, [0]] > 128\n",
    "    train_data = np.transpose(train_data.astype('uint8'), (0, 3, 1, 2))\n",
    "    test_data = np.transpose(test_data.astype('uint8'), (0, 3, 1, 2))\n",
    "    if flatten:\n",
    "        train_data = train_data.reshape(-1, 28 * 28)\n",
    "        test_data = test_data.reshape(-1, 28 * 28)\n",
    "    return train_data, test_data\n",
    "\n",
    "\n",
    "def show_samples(samples, title, nrow=10):\n",
    "    samples = torch.FloatTensor(samples).reshape(-1, 28, 28)\n",
    "    samples = torch.unsqueeze(samples, axis=1)\n",
    "    grid_img = make_grid(samples, nrow=nrow)\n",
    "    plt.figure()\n",
    "    plt.title(title)\n",
    "    plt.imshow(grid_img.permute(1, 2, 0))\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_mnist_images(data, title):\n",
    "    idxs = np.random.choice(len(data), replace=False, size=(100,))\n",
    "    images = train_data[idxs]\n",
    "    show_samples(images, title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DG0ctG8dxSIe"
   },
   "source": [
    "Here are the functions that we will you for training our model. Please, explore these functions carefully. You do not have to change them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gvEMBEn3r1z0"
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, train_loader, optimizer, use_cuda, loss_key='total'):\n",
    "    model.train()\n",
    "  \n",
    "    stats = defaultdict(list)\n",
    "    for x in train_loader:\n",
    "        if use_cuda:\n",
    "            x = x.cuda()\n",
    "        losses = model.loss(x)\n",
    "        optimizer.zero_grad()\n",
    "        losses[loss_key].backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        for k, v in losses.items():\n",
    "            stats[k].append(v.item())\n",
    "    return stats\n",
    "\n",
    "\n",
    "def eval_model(model, data_loader, use_cuda):\n",
    "    model.eval()\n",
    "    stats = defaultdict(float)\n",
    "    with torch.no_grad():\n",
    "        for x in data_loader:\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "            losses = model.loss(x)\n",
    "            for k, v in losses.items():\n",
    "                stats[k] += v.item() * x.shape[0]\n",
    "\n",
    "        for k in stats.keys():\n",
    "            stats[k] /= len(data_loader.dataset)\n",
    "    return stats\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, test_loader, epochs, lr, use_tqdm=False, use_cuda=False, loss_key='total_loss'):\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "    train_losses = defaultdict(list)\n",
    "    test_losses = defaultdict(list)\n",
    "    forrange = tqdm(range(epochs)) if use_tqdm else range(epochs)\n",
    "    if use_cuda:\n",
    "        model = model.cuda()\n",
    "    for epoch in forrange:\n",
    "        model.train()\n",
    "        train_loss = train_epoch(model, train_loader, optimizer, use_cuda, loss_key)\n",
    "        test_loss = eval_model(model, test_loader, use_cuda)\n",
    "\n",
    "        for k in train_loss.keys():\n",
    "            train_losses[k].extend(train_loss[k])\n",
    "            test_losses[k].append(test_loss[k])\n",
    "    return dict(train_losses), dict(test_losses)\n",
    "\n",
    "\n",
    "def plot_training_curves(train_losses, test_losses):\n",
    "    n_train = len(train_losses[list(train_losses.keys())[0]])\n",
    "    n_test = len(test_losses[list(train_losses.keys())[0]])\n",
    "    x_train = np.linspace(0, n_test - 1, n_train)\n",
    "    x_test = np.arange(n_test)\n",
    "\n",
    "    plt.figure()\n",
    "    for key, value in train_losses.items():\n",
    "        plt.plot(x_train, value, label=key + '_train')\n",
    "\n",
    "    for key, value in test_losses.items():\n",
    "        plt.plot(x_test, value, label=key + '_test')\n",
    "\n",
    "    plt.legend(fontsize=12)\n",
    "    plt.xlabel('Epoch', fontsize=14)\n",
    "    plt.ylabel('Loss', fontsize=14)\n",
    "    plt.xticks(fontsize=12)\n",
    "    plt.yticks(fontsize=12)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XnpPa7n7P7NA"
   },
   "source": [
    "## Task 1: PixelCNN on MNIST (5pt)\n",
    "\n",
    "[PixelCNN](https://arxiv.org/abs/1601.06759) model uses masked causal convoultions on images, we have discussed this model on lecture 2.\n",
    "\n",
    "Here you have to train this model on MNIST images.\n",
    "\n",
    "(see paper for details: https://arxiv.org/abs/1601.06759)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "7RN1T58NP-0Z",
    "outputId": "7e33128b-c14c-405f-9213-d534c9219cec"
   },
   "outputs": [],
   "source": [
    "# change the path to the file\n",
    "train_data, test_data = load_pickle(\n",
    "    os.path.join('drive', 'My Drive', 'DGM', 'homework_supplementary', 'mnist.pkl'),\n",
    "    flatten=False\n",
    ")\n",
    "visualize_mnist_images(train_data, 'MNIST samples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b3KCV1k7QAdS"
   },
   "outputs": [],
   "source": [
    "class MaskedConv2d(nn.Conv2d):\n",
    "    def __init__(self, mask_type, in_channels, out_channels, kernel_size=5):\n",
    "        assert mask_type in ['A', 'B']\n",
    "        super().__init__(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=kernel_size // 2)\n",
    "        self.register_buffer('mask', torch.zeros_like(self.weight))\n",
    "        self.create_mask(mask_type)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # ====\n",
    "        # your code\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def create_mask(self, mask_type):\n",
    "        # ====\n",
    "        # your code\n",
    "        # do not forget about mask_type\n",
    "        \n",
    "        # ====\n",
    "\n",
    "\n",
    "def test_masked_conv2d():\n",
    "    layer = MaskedConv2d('A', 2, 2)\n",
    "    assert np.allclose(layer.mask[:, :, 2, 2].numpy(), np.zeros((2, 2)))\n",
    "\n",
    "    layer = MaskedConv2d('B', 2, 2)\n",
    "    assert np.allclose(layer.mask[:, :, 2, 2].numpy(), np.ones((2, 2)))\n",
    "\n",
    "\n",
    "test_masked_conv2d()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LIsaX7QHrem1"
   },
   "source": [
    "[Layer Normalization](https://arxiv.org/abs/1607.06450) helps to stabilize training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "svKD0W21QCC2"
   },
   "outputs": [],
   "source": [
    "class LayerNorm(nn.LayerNorm):\n",
    "    def __init__(self, n_filters):\n",
    "        super().__init__(n_filters)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 2, 3, 1).contiguous()\n",
    "        x = super().forward(x)\n",
    "        return x.permute(0, 3, 1, 2).contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HGLaZzD4QEim"
   },
   "outputs": [],
   "source": [
    "class PixelCNN(nn.Module):\n",
    "    def __init__(\n",
    "        self, \n",
    "        input_shape, \n",
    "        n_filters=64, \n",
    "        kernel_size=7, \n",
    "        n_layers=5, \n",
    "        use_layer_norm=True\n",
    "    ):\n",
    "      \n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # apply the sequence of MaskedConv2d -> LayerNorm -> ReLU\n",
    "        # note that the first conv layer should be of type 'A'\n",
    "        # the last layer should be MaskedConv2d\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        out = (x.float() - 0.5) / 0.5\n",
    "        out = self.net(out)\n",
    "        return out.view(batch_size, 2, 1, *self.input_shape)\n",
    "\n",
    "    def loss(self, x):\n",
    "        # ====\n",
    "        # your code\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def sample(self, n):\n",
    "        # read carefully the sampling process\n",
    "        samples = torch.zeros(n, 1, *self.input_shape).cuda()\n",
    "        with torch.no_grad():\n",
    "            for r in range(self.input_shape[0]):\n",
    "                for c in range(self.input_shape[1]):\n",
    "                    logits = self(samples)[:, :, :, r, c]\n",
    "                    probs = F.softmax(logits, dim=1).squeeze(-1)\n",
    "                    samples[:, 0, r, c] = torch.multinomial(probs, num_samples=1).squeeze(-1)\n",
    "        return samples.permute(0, 2, 3, 1).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cje3w0aJQEfW",
    "outputId": "3000eda8-8b52-4927-b8ff-27f7e353f7bb"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "EPOCHS = \n",
    "BATCH_SIZE = \n",
    "LR = \n",
    "N_LAYERS = \n",
    "N_FILTERS = \n",
    "# ====\n",
    "\n",
    "model = PixelCNN(\n",
    "    input_shape=(28, 28), \n",
    "    n_filters=N_FILTERS, \n",
    "    kernel_size=5, \n",
    "    n_layers=N_LAYERS, \n",
    "    use_layer_norm=True\n",
    ")\n",
    "\n",
    "loss = model.loss(torch.zeros(1, 1, 28, 28))\n",
    "assert isinstance(loss, dict)\n",
    "assert 'total_loss' in loss\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "train_losses, test_losses = train_model(model, train_loader, test_loader, epochs=EPOCHS, lr=LR, use_tqdm=True, use_cuda=USE_CUDA)\n",
    "\n",
    "assert test_losses['total_loss'][-1] < 0.11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W4ftGoIzQI8b"
   },
   "outputs": [],
   "source": [
    "plot_training_curves(train_losses, test_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zac6FH7YQI_E"
   },
   "outputs": [],
   "source": [
    "samples = model.sample(25)\n",
    "show_samples(samples, title='MNIST samples', nrow=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2SmVPIL-QMc3"
   },
   "source": [
    "The PixelCNN model is a powerful model. The samples should have better quality compared to MADE model from HW1. But the model has drawbacks.\n",
    "\n",
    "1. The model is sequential and sampling is really slow (it is a drawback of all AR models).\n",
    "\n",
    "2. The receptive field of the model is not so large. Even if the model is well-trained, the samples do not have long-range history. To validate this let try to visualize the receptive field of the model. In order to do this we will use the following function. It is really helpful to understand this code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2G-fpPq3QOPI"
   },
   "outputs": [],
   "source": [
    "def plot_receptive_field(model, data, n_layers):\n",
    "    out = model(data)\n",
    "    x_center, y_center = data.shape[-2] // 2, data.shape[-1] // 2\n",
    "    out[0, 0, 0, x_center, y_center].backward()\n",
    "    grad = data.grad.detach().cpu().numpy()[0, 0]\n",
    "    grad = np.abs(grad)\n",
    "    binary_map = (grad > 1e-8).astype('float32')\n",
    "    binary_map = np.stack([binary_map, binary_map, binary_map], axis=-1)\n",
    "    binary_map[x_center, y_center] = [1, 0, 0]\n",
    "    weighted_map = grad * (1 / grad.max())\n",
    "    weighted_map = np.stack([weighted_map, weighted_map, weighted_map], axis=-1)\n",
    "    weighted_map[x_center, y_center] = [1, 0, 0]\n",
    "\n",
    "    fig, ax = plt.subplots(1,2, figsize=(10, 6))\n",
    "    ax[0].imshow(weighted_map, vmin=0.0, vmax=1.0)\n",
    "    ax[1].imshow(binary_map, vmin=0.0, vmax=1.0)\n",
    "\n",
    "    ax[0].set_title(f\"Weighted receptive field, PixelCNN {n_layers} layers\")\n",
    "    ax[1].set_title(f\"Binary receptive field, PixelCNN {n_layers} layers\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bf-yRHo_QOR2"
   },
   "outputs": [],
   "source": [
    "x = torch.zeros(1, 1, 28, 28)\n",
    "\n",
    "if USE_CUDA:\n",
    "    x = x.cuda()\n",
    "x.requires_grad = True\n",
    "\n",
    "for n_layers in [1, 3, 5, 6]:\n",
    "    model = PixelCNN(\n",
    "        input_shape=(28, 28), \n",
    "        n_filters=32, \n",
    "        kernel_size=5, \n",
    "        n_layers=n_layers, \n",
    "        use_layer_norm=True\n",
    "    )\n",
    "    if USE_CUDA:\n",
    "        model = model.cuda()\n",
    "    plot_receptive_field(model, x, n_layers)\n",
    "    x.grad.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P60O_7hE9isj"
   },
   "source": [
    "You have to see that PixelCNN has strange blind spot in binary receptive field plot on the right side. This is a known issue of PixelCNN model. Please, try to understand why it happens. \n",
    "\n",
    "One way to solve this problem is a GatedPixelCNN model (see paper https://arxiv.org/pdf/1606.05328.pdf, if you are interested in)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zGamQ3L4W64G"
   },
   "source": [
    "## Task 2: Theory (3pt)\n",
    "\n",
    "In Lecture 3, when deriving the ELBO gradient at the E-step, we encountered a problem with Monte Carlo estimation, since the distribution function depended on the differentiation parameters.\n",
    "\\begin{align*}\n",
    "    \\nabla_{\\boldsymbol{\\phi}} \\mathcal{L} (\\boldsymbol{\\phi}, \\boldsymbol{\\theta}) &= \\nabla_{\\boldsymbol{\\phi}} \\int q(\\mathbf{z} | \\mathbf{x}, \\boldsymbol{\\phi}) \\left[\\log p(\\mathbf{x}, \\mathbf{z} | \\boldsymbol{\\theta}) - \\log q(\\mathbf{z}| \\mathbf{x}, \\boldsymbol{\\phi}) \\right] d \\mathbf{z} \\\\\n",
    "    & \\neq  \\int q(\\mathbf{z} | \\mathbf{x}, \\boldsymbol{\\phi}) \\nabla_{\\boldsymbol{\\phi}} \\left[\\log p(\\mathbf{x}, \\mathbf{z} | \\boldsymbol{\\theta}) - \\log q(\\mathbf{z}| \\mathbf{x}, \\boldsymbol{\\phi}) \\right] d \\mathbf{z} \\\\\n",
    "\\end{align*}\n",
    "The Reparametrization trick allowed us to skip the gradient and get a Monte Carlo estimate. But there is another way that uses the so-called \\textbf{log-derivative trick}:\n",
    "$$\n",
    "    \\nabla_\\xi  \\log q(\\eta| \\xi) = \\frac{\\nabla_\\xi q(\\eta| \\xi)}{q(\\eta| \\xi)}.\n",
    "$$\n",
    "* Get the Monte Carlo estimate of the gradient using the formula for the derivative of the logarithm.\n",
    "* The resulting expression works significantly worse than the reparametrization trick. Namely, it has a huge variance. Try to describe the intuition why the evaluation has a high variance (you need to think about what order and sign the terms in the expression will have)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-7SKI68oY2OF"
   },
   "source": [
    "```\n",
    "your solution\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aZjBDpnudlJL"
   },
   "source": [
    "## Task 2: VAE on 2d data (5pt)\n",
    "\n",
    "In this task you will implement simple VAE model for 2d gaussian distribution.\n",
    "\n",
    "We will consider two cases: 2d univariate distribution and 2d multivariate distribution. The goal is to analyze the difference between these two cases and chosen VAE model.\n",
    "\n",
    "Below you will find data generation function. Look carefully, do not change."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PksKzhF733Np"
   },
   "source": [
    "Below you will find the functions which visualizes 2d data and plots the obtained training curves.\n",
    "Study them carefully but do not change them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zdTy9DLja_jV"
   },
   "outputs": [],
   "source": [
    "def generate_2d_data(count, mode='univariate'):\n",
    "    assert mode in ['univariate', 'multivariate']\n",
    "    np.random.seed(42)\n",
    "    mean = [[2.0, 3.0]]\n",
    "    sigma = [[3.0, 1.0]]\n",
    "    if mode == 'univariate':\n",
    "        rotate = [\n",
    "            [1.0, 0.0], \n",
    "            [0.0, 1.0]\n",
    "        ]\n",
    "    else:\n",
    "        rotate = [\n",
    "            [np.sqrt(2) / 2, np.sqrt(2) / 2], \n",
    "            [-np.sqrt(2) / 2, np.sqrt(2) / 2]\n",
    "        ]\n",
    "    data = mean + (np.random.randn(count, 2) * sigma).dot(rotate)\n",
    "    data = data.astype('float32')\n",
    "    split = int(0.7 * count)\n",
    "    train_data, test_data = data[:split], data[split:]\n",
    "    return train_data, test_data\n",
    "\n",
    "\n",
    "def visualize_2d_data(train_data, test_data, train_labels=None, test_labels=None):\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))\n",
    "    ax1.set_title('train', fontsize=16)\n",
    "    ax1.scatter(train_data[:, 0], train_data[:, 1], s=1, c=train_labels)\n",
    "    ax1.tick_params(labelsize=16)\n",
    "    ax2.set_title('test', fontsize=16)\n",
    "    ax2.scatter(test_data[:, 0], test_data[:, 1], s=1, c=test_labels)\n",
    "    ax2.tick_params(labelsize=16)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_2d_samples(data, title):\n",
    "    plt.figure(figsize=(5, 5))\n",
    "    plt.scatter(data[:, 0], data[:, 1], s=1)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def get_latent_stats(model, test_data, batch_size=3000):\n",
    "    batch = next(iter(data.DataLoader(test_data, batch_size=batch_size, shuffle=True)))\n",
    "    batch = batch.cuda()\n",
    "    with torch.no_grad():\n",
    "        \n",
    "        mu_z, log_std_z = model(batch)[:2]\n",
    "        \n",
    "    mu_z = mu_z.cpu().numpy()\n",
    "    std_z = log_std_z.exp().cpu().numpy()\n",
    "\n",
    "    return mu_z, std_z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JlZlxRm3a8Mi"
   },
   "outputs": [],
   "source": [
    "COUNT = 15000\n",
    "\n",
    "train_data, test_data = generate_2d_data(COUNT, mode='multivariate')\n",
    "visualize_2d_data(train_data, test_data)\n",
    "\n",
    "train_data, test_data = generate_2d_data(COUNT, mode='univariate')\n",
    "visualize_2d_data(train_data, test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vRYd_B3f3jz"
   },
   "source": [
    "The difference of these two cases is the form of covariance matrix.\n",
    "\n",
    "In multivariate case the matrix is non-diagonal, in univariate case it is strictly diagonal. As you will see, our VAE model will have absolutely different results for these datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "biYy9_rWd-DY"
   },
   "source": [
    "Now it is time to define our model. In this task you will create VAE model on 2d data. Model will be designed as follows:\n",
    "\n",
    "* the latent dimensionality is equal to 2, the same as the data dimensionality ($\\mathbf{z} \\in \\mathbb{R}^2$, $\\mathbf{x} \\in \\mathbb{R}^2$).\n",
    "* prior distribution $p(\\mathbf{z}) = \\mathcal{N}(0, I)$.\n",
    "* approximate variational distribution (or encoder) $q(\\mathbf{z} | \\mathbf{x}, \\boldsymbol{\\phi}) = \\mathcal{N}(\\boldsymbol{\\mu}_{\\boldsymbol{\\phi}}(\\mathbf{x}), \\boldsymbol{\\Sigma}_{\\boldsymbol{\\phi}}(\\mathbf{x}))$. Here $\\boldsymbol{\\phi}$ denotes all parameters of the encoder neural network. \n",
    "* generative distribution (or decoder) $p(\\mathbf{x} | \\mathbf{z}, \\boldsymbol{\\theta}) = \\mathcal{N}(\\boldsymbol{\\mu}_{\\boldsymbol{\\theta}}(\\mathbf{z}), \\boldsymbol{\\Sigma}_{\\boldsymbol{\\theta}}(\\mathbf{z}))$. Here $\\boldsymbol{\\theta}$ denotes all parameters of the decoder neural network. Please note, that here we will use continuous distribution at our variables.\n",
    "* We will consider only diagonal covariance matrices $\\boldsymbol{\\Sigma}_{\\boldsymbol{\\phi}}(\\mathbf{x})$, $\\boldsymbol{\\Sigma}_{\\boldsymbol{\\theta}}(\\mathbf{z})$.\n",
    "\n",
    "Model objective is ELBO:\n",
    "$$\n",
    "    \\mathcal{L}(\\boldsymbol{\\phi}, \\boldsymbol{\\theta}) = \\mathbb{E}_{q(\\mathbf{z} | \\mathbf{x}, \\boldsymbol{\\phi})} \\log p(\\mathbf{x} | \\mathbf{z}, \\boldsymbol{\\theta}) - KL (q(\\mathbf{z} | \\mathbf{x}, \\boldsymbol{\\phi}) || p(\\mathbf{z})).\n",
    "$$\n",
    "\n",
    "To make the expectation is independent of parameters $\\boldsymbol{\\phi}$, we will use reparametrization trick.\n",
    "\n",
    "To calculate the loss, you should derive\n",
    "- $\\log p(\\mathbf{x} | \\mathbf{z}, \\boldsymbol{\\theta})$, note that generative distribution is $\\mathcal{N}(\\boldsymbol{\\mu}_{\\boldsymbol{\\theta}}(\\mathbf{z}), \\boldsymbol{\\Sigma}_{\\boldsymbol{\\theta}}(\\mathbf{z}))$.\n",
    "- KL between $\\mathcal{N}(\\boldsymbol{\\mu}_{\\boldsymbol{\\phi}}(\\mathbf{x}), \\boldsymbol{\\Sigma}_{\\boldsymbol{\\phi}}(\\mathbf{x}))$ and $\\mathcal{N}(0, I)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_ZEUyHuxE39I"
   },
   "outputs": [],
   "source": [
    "def get_normal_KL(mean_1, log_std_1, mean_2=None, log_std_2=None):\n",
    "    \"\"\"\n",
    "        This function should return the value of KL(p1 || p2),\n",
    "        where p1 = Normal(mean_1, exp(log_std_1)), p2 = Normal(mean_2, exp(log_std_2) ** 2).\n",
    "        If mean_2 and log_std_2 are None values, we will use standart normal distribution.\n",
    "        Note that we consider the case of diagonal covariance matrix.\n",
    "    \"\"\"\n",
    "    if mean_2 is None:\n",
    "        mean_2 = torch.zeros_like(mean_1)\n",
    "    if log_std_2 is None:\n",
    "        log_std_2 = torch.zeros_like(log_std_1)\n",
    "    # ====\n",
    "    # your code\n",
    "    \n",
    "    # ====\n",
    "\n",
    "\n",
    "def test_KL():\n",
    "    assert np.isclose(get_normal_KL(torch.tensor(2), torch.tensor(3), torch.tensor(0), torch.tensor(0)).numpy(), 200.2144, rtol=1e-3)\n",
    "    assert np.isclose(get_normal_KL(torch.tensor(2), torch.tensor(3), torch.tensor(4), torch.tensor(5)).numpy(), 1.50925, rtol=1e-3)\n",
    "    assert np.allclose(get_normal_KL(torch.tensor((10, 10)), torch.tensor((2, 4)), torch.tensor((3, 5))).numpy(), [49.2990, 1498.479], rtol=1e-3)\n",
    "\n",
    "test_KL()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "871Pfpm1TiWF"
   },
   "outputs": [],
   "source": [
    "def get_normal_nll(x, mean, log_std):\n",
    "    \"\"\"\n",
    "        This function should return the negative log likelihood log p(x),\n",
    "        where p(x) = Normal(x | mean, exp(log_std) ** 2).\n",
    "        Note that we consider the case of diagonal covariance matrix.\n",
    "    \"\"\"\n",
    "    # ====\n",
    "    # your code\n",
    "    \n",
    "    # ====\n",
    "\n",
    "\n",
    "def test_NLL():\n",
    "    assert np.isclose(get_normal_nll(torch.tensor(2), torch.tensor(2), torch.tensor(3)).numpy(), 3.9189, rtol=1e-3)\n",
    "    assert np.isclose(get_normal_nll(torch.tensor(5), torch.tensor(-3), torch.tensor(6)).numpy(), 6.9191, rtol=1e-3)\n",
    "    assert np.allclose(get_normal_nll(torch.tensor((10, 10)), torch.tensor((2, 4)), torch.tensor((3, 5))).numpy(), np.array([3.9982, 5.9197]), rtol=1e-3)\n",
    "\n",
    "test_NLL()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KN7NnFplkUSn"
   },
   "outputs": [],
   "source": [
    "class FullyConnectedMLP(nn.Module):\n",
    "    def __init__(self, input_shape, hiddens, output_shape):\n",
    "        assert isinstance(hiddens, list)\n",
    "        super().__init__()\n",
    "        self.input_shape = (input_shape,)\n",
    "        self.output_shape = (output_shape,)\n",
    "        self.hiddens = hiddens\n",
    "\n",
    "        # ====\n",
    "        # your code \n",
    "        # Stack Dense layers with ReLU activation.\n",
    "        # Note that you do not have to add relu after the last dense layer\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def forward(self, x):\n",
    "        # ====\n",
    "        # your code\n",
    "        # apply network that was defined in __init__ and return the output\n",
    "        \n",
    "        # ====\n",
    "\n",
    "\n",
    "class VAE2d(nn.Module):\n",
    "    def __init__(self, n_in, n_latent, enc_hidden_sizes, dec_hidden_sizes):\n",
    "        assert isinstance(enc_hidden_sizes, list)\n",
    "        assert isinstance(dec_hidden_sizes, list)\n",
    "        super().__init__()\n",
    "        self.n_latent = n_latent\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # define encoder and decoder networks\n",
    "        # the encoder takes n_in elements, has enc_hidden_sizes neurons in hidden layers \n",
    "        # and outputs 2 * n_latent (n_latent for means, and n_latent for std)\n",
    "        # the decoder takes n_latent elements, has dec_hidden_sizes neurons in hidden layers \n",
    "        # and outputs 2 * n_in (n_in for means, and n_in for std)\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def prior(self, n):\n",
    "        # ====\n",
    "        # your code\n",
    "        # return n samples from prior distribution (we use standart normal for prior)\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def forward(self, x):\n",
    "        # ====\n",
    "        # your code\n",
    "        # now you have to return from the model \n",
    "        # - mu_z - means for variational distribution \n",
    "        # - mu_x - means for generative distribution\n",
    "        # - log_std_z - logarithm of std for variational distribution\n",
    "        # - log_std_x - logarithm of std for generative distribution\n",
    "        # we use logarithm, since the std is always positive\n",
    "        # to get std we will exponentiate it to get rid of this constraint\n",
    "        # 1) mu_z, log_std_z are outputs from the encoder\n",
    "        # 2) apply reparametrization trick to get z (input of decoder)\n",
    "        # (do not forget to use self.prior())\n",
    "        # 3) mu_x, log_std_x are outputs from the decoder\n",
    "        \n",
    "        # ====\n",
    "        return mu_z, log_std_z, mu_x, log_std_x\n",
    "\n",
    "    def loss(self, x):\n",
    "        mu_z, log_std_z, mu_x, log_std_x = self(x)\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply model to get mu_z, log_std_z, mu_x, log_std_x\n",
    "        # 2) compute reconstruction loss using get_normal_nll (it is the first term in ELBO)\n",
    "        # 3) compute KL loss using get_normal_KL (it is the second term in ELBO)\n",
    "        \n",
    "        # ====\n",
    "\n",
    "        return {\n",
    "            'elbo_loss': recon_loss + kl_loss, \n",
    "            'recon_loss': recon_loss,\n",
    "            'kl_loss': kl_loss\n",
    "        }\n",
    "\n",
    "    def sample(self, n, noise=True):\n",
    "        with torch.no_grad():\n",
    "            # ====\n",
    "            # your code\n",
    "            # to sample from VAE model you have to sample from prior\n",
    "            # and then apply decoder to prior samples.\n",
    "            # parameter noise indicates whether to sample from decoder\n",
    "            # or just use means of generative distribution as samples\n",
    "            # 1) generate prior samples\n",
    "            # 2) apply decoder\n",
    "            # 3) sample from the decoder distribution if noise=True\n",
    "            \n",
    "            # ====\n",
    "        return z.cpu().numpy()\n",
    "\n",
    "\n",
    "def solve_task(train_data, test_data, model, batch_size, epochs, lr, use_cuda=False, use_tqdm=False):\n",
    "    # do not change this function\n",
    "    train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "    test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "    train_losses, test_losses = train_model(\n",
    "        model, train_loader, test_loader, epochs=EPOCHS, lr=LR, use_cuda=use_cuda, use_tqdm=use_tqdm, loss_key='elbo_loss'\n",
    "    )\n",
    "    samples_noise = model.sample(3000, noise=True)\n",
    "    samples_nonoise = model.sample(3000, noise=False)\n",
    "\n",
    "    for key, value in test_losses.items():\n",
    "        print('{}: {:.4f}'.format(key, value[-1]))\n",
    "\n",
    "    plot_training_curves(train_losses, test_losses)\n",
    "    visualize_2d_samples(samples_noise, title='Samples with Decoder Noise')\n",
    "    visualize_2d_samples(samples_nonoise, title='Samples without Decoder Noise')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fDAWwEs8eJWV"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters (2 hidden layers could be enough for encoder and decoder)\n",
    "ENC_HIDDEN_SIZES = \n",
    "DEC_HIDDEN_SIZES = \n",
    "BATCH_SIZE = \n",
    "EPOCHS = \n",
    "LR = \n",
    "# ====\n",
    "\n",
    "COUNT = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IdJaAD6Ls4hL"
   },
   "outputs": [],
   "source": [
    "train_data, test_data = generate_2d_data(COUNT, mode='multivariate')\n",
    "visualize_2d_data(train_data, test_data)\n",
    "\n",
    "model = VAE2d(2, 2, ENC_HIDDEN_SIZES, DEC_HIDDEN_SIZES).cuda()\n",
    "solve_task(train_data, test_data, model, BATCH_SIZE, EPOCHS, LR, use_cuda=USE_CUDA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FbtO7n9GOsc2"
   },
   "outputs": [],
   "source": [
    "# just look at these numbers and read comments after this task\n",
    "mu_z, std_z = get_latent_stats(model, test_data)\n",
    "\n",
    "print('mu_z = ', mu_z.mean(axis=0), '+-', mu_z.std(axis=0))\n",
    "print('std_z = ', std_z.mean(axis=0), '+-', std_z.std(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ob-yRJ8xe7Ns"
   },
   "outputs": [],
   "source": [
    "train_data, test_data = generate_2d_data(COUNT, mode='univariate')\n",
    "visualize_2d_data(train_data, test_data)\n",
    "\n",
    "model = VAE2d(2, 2, ENC_HIDDEN_SIZES, DEC_HIDDEN_SIZES).cuda()\n",
    "solve_task(train_data, test_data, model, BATCH_SIZE, EPOCHS, LR, use_cuda=USE_CUDA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KljEJJWdNyx-"
   },
   "outputs": [],
   "source": [
    "# just look at these numbers and read comments after this task\n",
    "mu_z, std_z = get_latent_stats(model, test_data)\n",
    "\n",
    "print('mu_z = ', mu_z.mean(axis=0), '+-', mu_z.std(axis=0))\n",
    "print('std_z = ', std_z.mean(axis=0), '+-', std_z.std(axis=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lRCmH4X_6tyQ"
   },
   "source": [
    "After training the model on these 2 datasets, have a look at \"Samples without Decoder Noise\" figures. These figures show the means of the generative distribution. In the case of multivariate gaussian, the means are perfectly aligned with data distribution. \n",
    "Otherwise, in the univariate gaussian case you have to see strange figure. This happens due to so called \"posterior collapse\" (we will discuss it on the one of our lecture).\n",
    "\n",
    "\n",
    "Our posterior distribution $p(\\mathbf{x} | \\mathbf{z}, \\boldsymbol{\\theta}) = \\mathcal{N}(\\boldsymbol{\\mu}_{\\boldsymbol{\\theta}}(\\mathbf{z}), \\boldsymbol{\\Sigma}_{\\boldsymbol{\\theta}}(\\mathbf{z}))$ is a univariate (covariance matrix is diagonal). Thus, to model univariate data distribution (second case) the model ignores latent code, cause the model fits the distribution without any information from latent space.\n",
    "\n",
    "If the decoder ignores latent code, the second term in ELBO (KL) could be low (variational distribution, which is given by encoder model, is close to prior distribution for each datapoint). In the training curves you have to see that KL loss behaves differently in these two cases.\n",
    "\n",
    "The mean and std of variational distribution also proves this concept. For the second case you have to see that mean is almost zero and std is almost one.\n",
    "\n",
    "It is a real problem for generative models and we will discuss later how to overcome it."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "hw2.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}