homeworks/hw1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5rUp2xWkdGFu"
   },
   "source": [
    "# Homework 1: Autoregressive models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tk8-eCJuaTrU"
   },
   "source": [
    "## Task 1. Theory (5pt)\n",
    "\n",
    "1. Consider the MADE model with a single hidden layer. The input object is $\\mathbf{x} \\in \\mathbb{R}^m$. We denote by $\\mathbf{W} \\in \\mathbb{R}^{h \\times m}$ the matrix of weights between the input and the hidden layer, and by $\\mathbf{V} \\in \\mathbb{R}^{m \\times h}$ the matrix of weights between the hidden and the output layer ($h$ is the number of neurons in the hidden layer). Let us generate the correct autoregressive masks $\\mathbf{M}_{\\mathbf{W}} \\in \\mathbb{R}^{h \\times m}$ and $\\mathbf{M}_{\\mathbf{V}} \\in \\mathbb{R}^{m \\times h}$ (the generation algorithm is given in Lecture 1) for the direct order of variables\n",
    "$$\n",
    "    p(\\mathbf{x}) = p(x_1) \\cdot p(x_2 | x_1) \\cdot \\dots \\cdot p(x_m | x_{m-1}, \\dots, x_1).\n",
    "$$ \n",
    "Each mask is a binary matrix of 0 and 1. Let's introduce the matrix $\\mathbf{M} = \\mathbf{M}_{\\mathbf{V}} \\mathbf{M}_{\\mathbf{W}}$. Prove that:\n",
    "    * $\\mathbf{M}$ is strictly lower triangular (has zeros on the diagonal and above the diagonal);\n",
    "    * $\\mathbf{M}_{ij}$  is equal to the number of paths in the network graph between the output neuron $\\hat{x}_i$ and the input neuron $x_j$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "erVRXwzqPHwV"
   },
   "source": [
    "```\n",
    "your solution for task 1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QLGp4c5UPByO"
   },
   "source": [
    "\n",
    "2. Let's suppose we have 2 generative models for images of size $W \\times H \\times C$, where $W$ - image width, $H$ - image height, $C$ - number of channels. The first model $p_1(\\mathbf{x} | \\boldsymbol{\\theta})$ outputs a discrete distribution for each pixel  $\\text{Categorical}(\\boldsymbol{\\pi})$, где $\\boldsymbol{\\pi} = (\\pi_1, \\dots,  \\pi_{256})$. The second model $p_2(\\mathbf{x} | \\boldsymbol{\\theta})$ models a discrete distribution by a continuous mixture of logistic functions\n",
    "$$\n",
    "    p(\\nu | \\boldsymbol{\\mu}, \\mathbf{s}, \\boldsymbol{\\pi}) = \\sum_{k=1}^K \\pi_k p(\\nu | \\mu_k, s_k).\n",
    "    $$\n",
    "$$\n",
    "    P(x | \\boldsymbol{\\mu}, \\mathbf{s}, \\boldsymbol{\\pi}) = P(x + 0.5 | \\boldsymbol{\\mu}, \\mathbf{s}, \\boldsymbol{\\pi}) - P(x - 0.5 | \\boldsymbol{\\mu}, \\mathbf{s}, \\boldsymbol{\\pi}).\n",
    "$$\n",
    "    Each of the models outputs parameters of pixel distributions.\n",
    "\n",
    "    * Calculate the dimensions of the output tensor for the model $p_1(\\mathbf{x} | \\boldsymbol{\\theta})$ and for the model $p_2(\\mathbf{x} | \\boldsymbol{\\theta})$. \n",
    "    * At what number of mixture components $K$ is the number of elements of the output tensor for $p_2(\\mathbf{x} | \\boldsymbol{\\theta})$ becomes greater than $p_1(\\mathbf{x} | \\boldsymbol{\\theta})$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P0HFuytJPQ7o"
   },
   "source": [
    "```\n",
    "your solution for task 2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ivV4DY1SPEIL"
   },
   "source": [
    "3. In the course, we will meet different divergences (not only $KL$). So let's get acquainted with the class of $\\alpha$ - divergences:\n",
    "$$\n",
    "    D_{\\alpha}(p || q) = \\frac{4}{1 - \\alpha^2} \\left( 1 - \\int p(x)^{\\frac{1 + \\alpha}{2}}q(x)^{\\frac{1 - \\alpha}{2}}dx\\right).\n",
    "$$\n",
    "For each $\\alpha \\in [-\\infty; +\\infty]$ the function $D_{\\alpha} (p || q)$ is a measure of the similarity of the two distributions, which could have different properties.\n",
    "\t  \n",
    "      Prove that for $\\alpha \\rightarrow 1$ the divergence $D_{\\alpha}(p || q) \\rightarrow KL(p || q)$, and for $\\alpha \\rightarrow -1$ the divergence $D_{\\alpha}(p || q) \\rightarrow KL(q || p)$. (Hint: use the fact that $t^\\epsilon = \\exp(\\epsilon \\ln t) = 1 + \\epsilon \\ln t + O(\\epsilon^2)$.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0UP5Y6Z4PSP3"
   },
   "source": [
    "```\n",
    "your solution for task 3\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d2LH7vKPPVUt"
   },
   "source": [
    "Now it time to move on to practical part of homework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zk6rWePvdGFv"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import io\n",
    "import itertools\n",
    "import pickle\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_m5NVPFaJGHO",
    "outputId": "bd7590eb-29fc-48ab-d57d-3d1547e7f0e3"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K-wuVhhNdGFz"
   },
   "source": [
    "Use the following functions to train your models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U3MMoLe6dGFz"
   },
   "outputs": [],
   "source": [
    "def get_cross_entropy_loss(scores, labels):\n",
    "    # ====\n",
    "    # your code\n",
    "    \n",
    "    # ====\n",
    "\n",
    "\n",
    "def test_get_cross_entropy_loss():\n",
    "    input = torch.tensor([[1, 2, 3, 4],[5, 6, 7, 8]], dtype=torch.float32)\n",
    "    target = torch.tensor([3, 1], dtype=torch.long)\n",
    "\n",
    "    assert np.allclose(get_cross_entropy_loss(input, target).numpy(), 1.4402)\n",
    "\n",
    "test_get_cross_entropy_loss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GYg9WpssK66A"
   },
   "source": [
    "Do not change these functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3XzPpaW7JbLn"
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, train_loader, optimizer, use_cuda):\n",
    "    model.train()\n",
    "  \n",
    "    train_losses = []\n",
    "    for x in train_loader:\n",
    "        if use_cuda:\n",
    "            x = x.cuda()\n",
    "        loss = get_cross_entropy_loss(model(x), x)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_losses.append(loss.item())\n",
    "    return train_losses\n",
    "\n",
    "\n",
    "def eval_model(model, data_loader, use_cuda):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for x in data_loader:\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "            loss = get_cross_entropy_loss(model(x), x)\n",
    "            total_loss += loss * x.shape[0]\n",
    "        avg_loss = total_loss / len(data_loader.dataset)\n",
    "    return avg_loss.item()\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, test_loader, epochs, lr, use_tqdm=False, use_cuda=False):\n",
    "    if use_cuda:\n",
    "        model = model.cuda()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "    train_losses = []\n",
    "    test_losses = [eval_model(model, test_loader, use_cuda)]\n",
    "    if use_tqdm:\n",
    "        forrange = tqdm(range(epochs))\n",
    "    else:\n",
    "        forrange = range(epochs)\n",
    "    for epoch in forrange:\n",
    "        model.train()\n",
    "        train_losses.extend(train_epoch(model, train_loader, optimizer, use_cuda))\n",
    "        test_loss = eval_model(model, test_loader, use_cuda)\n",
    "        test_losses.append(test_loss)\n",
    "\n",
    "    return train_losses, test_losses\n",
    "\n",
    "\n",
    "def plot_training_curves(train_losses, test_losses):\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    n_epochs = len(test_losses) - 1\n",
    "    x_train = np.linspace(0, n_epochs, len(train_losses))\n",
    "    x_test = np.arange(n_epochs + 1)\n",
    "\n",
    "    plt.plot(x_train, train_losses, label='train loss')\n",
    "    plt.plot(x_test, test_losses, label='test loss')\n",
    "    plt.legend()\n",
    "    plt.title('training curves')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('NLL')\n",
    "\n",
    "\n",
    "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": "4GLIbwqkdGGD"
   },
   "source": [
    "## Task 2: MADE on 2D data (5pt)\n",
    "\n",
    "Train MADE model on single image (see paper for details: https://arxiv.org/abs/1502.03509).\n",
    "\n",
    "You will work with bivariate data of the form $x = (x_0,x_1)$, where $x_0, x_1 \\in \\{0, \\dots, \\text{n_bins}\\}$ (e.g. Categorial random variables). \n",
    "\n",
    "Implement and train a MADE model through MLE to represent $p(x_0, x_1)$ on the given image, with any autoregressive ordering of your choosing ($p(x_0, x_1) = p(x_0)p(x_1 | x_0)$ or $p(x_0, x_1) = p(x_1)p(x_0 | x_1)$). \n",
    "\n",
    "We advice you to think about what conditional distribution that you want to fit and how MADE's masks should look like. It may be useful to one-hot encode your inputs.\n",
    "\n",
    "You do not have to change these functions (except the path to the data file. Download the file from here: https://drive.google.com/file/d/1GUthJrA5fBpvi593Swo36t8zaFw9Dyak/view?usp=sharing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ngmqoNA2dGGD"
   },
   "outputs": [],
   "source": [
    "def generate_2d_data(count, bins):\n",
    "    # change the path to the image\n",
    "    im = Image.open(os.path.join('drive', 'My Drive', 'DGM', 'homework_supplementary', 'dgm.png')).resize((bins, bins)).convert('L')\n",
    "    im = np.array(im).astype('float32')\n",
    "    dist = im / im.sum()\n",
    "\n",
    "    pairs = list(itertools.product(range(bins), range(bins)))\n",
    "    idxs = np.random.choice(len(pairs), size=count, replace=True, p=dist.reshape(-1))\n",
    "    samples = np.array([pairs[i] for i in idxs])\n",
    "\n",
    "    split = int(0.8 * len(samples))\n",
    "    return dist, samples[:split], samples[split:]\n",
    "\n",
    "\n",
    "def plot_2d_data(train_data, test_data):\n",
    "    bins = int(max(test_data.max(), train_data.max()) - min(test_data.min(), train_data.min())) + 1\n",
    "    train_dist, test_dist = np.zeros((bins, bins)), np.zeros((bins, bins))\n",
    "    for i in range(len(train_data)):\n",
    "        train_dist[train_data[i][0], train_data[i][1]] += 1\n",
    "    train_dist /= train_dist.sum()\n",
    "\n",
    "    for i in range(len(test_data)):\n",
    "        test_dist[test_data[i][0], test_data[i][1]] += 1\n",
    "    test_dist /= test_dist.sum()\n",
    "\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 6))\n",
    "    ax1.set_title('Train Data')\n",
    "    ax1.imshow(train_dist, cmap='gray')\n",
    "    ax1.axis('off')\n",
    "    ax1.set_xlabel('x1')\n",
    "    ax1.set_ylabel('x0')\n",
    "\n",
    "    ax2.set_title('Test Data')\n",
    "    ax2.imshow(test_dist, cmap='gray')\n",
    "    ax2.axis('off')\n",
    "    ax2.set_xlabel('x1')\n",
    "    ax2.set_ylabel('x0')\n",
    "\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def plot_2d_distribution(true_dist, learned_dist):\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))\n",
    "    ax1.imshow(true_dist, cmap='gray')\n",
    "    ax1.set_title('True Distribution')\n",
    "    ax1.axis('off')\n",
    "    ax2.imshow(learned_dist, cmap='gray')\n",
    "    ax2.set_title('Learned Distribution')\n",
    "    ax2.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 301
    },
    "id": "eX5p_02jdGGG",
    "outputId": "6095558a-a560-4d90-e253-a9375962146b"
   },
   "outputs": [],
   "source": [
    "COUNT = 20000\n",
    "BINS = 60\n",
    "\n",
    "image, train_data, test_data = generate_2d_data(COUNT, BINS)\n",
    "plot_2d_data(train_data, test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dpcKxvPUMq27"
   },
   "source": [
    "Now we have to implement masked dense layer. It is a core component of MADE model. It acts like a usual dense layer, but firstly multiplies the weights by the predefined binary mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eLs_GDnMMoZg"
   },
   "outputs": [],
   "source": [
    "class MaskedLinear(nn.Linear):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super().__init__(in_features, out_features)\n",
    "        self.register_buffer('mask', torch.ones(out_features, in_features))\n",
    "\n",
    "    def set_mask(self, mask):\n",
    "        self.mask.data.copy_(torch.from_numpy(mask.astype(np.uint8).T))\n",
    "\n",
    "    def forward(self, input):\n",
    "        # ====\n",
    "        # your code\n",
    "        \n",
    "        # ====\n",
    "\n",
    "\n",
    "layer = MaskedLinear(2, 2)\n",
    "\n",
    "x = torch.tensor([1, 2], dtype=torch.float32)\n",
    "output = layer(x).detach().numpy()\n",
    "\n",
    "layer.set_mask(np.array([[0, 0], [0, 0]]))\n",
    "assert np.allclose(layer(x).detach().numpy(), layer.bias.detach().numpy())\n",
    "\n",
    "layer.set_mask(np.array([[1, 1], [1, 1]]))\n",
    "assert np.allclose(layer(x).detach().numpy(), output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R--RMmRgdGGI"
   },
   "outputs": [],
   "source": [
    "def to_one_hot(labels, d):\n",
    "    \"\"\"\n",
    "        The function takes categorical labels of size: batch_size x n_dims.\n",
    "        One-hot encodes them to d bins and then reshapes the result to batch_size x (n_dims * d)\n",
    "    \"\"\"\n",
    "    assert len(labels.shape) == 2\n",
    "    one_hot = F.one_hot(labels.to(torch.int64), d)\n",
    "    return one_hot.view((labels.shape[0], -1)).float()\n",
    "\n",
    "    \n",
    "class MADE(nn.Module):\n",
    "    def __init__(self, nin, bins, hidden_sizes):\n",
    "        super().__init__()\n",
    "        self.nin = nin\n",
    "        self.nout = nin * bins\n",
    "        self.bins = bins\n",
    "        self.hidden_sizes = hidden_sizes\n",
    "        # we will use the trivial ordering of input units\n",
    "        self.ordering = np.arange(self.nin)\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # define a simple MLP (sequence of MaskedLinear and ReLU) neural net \n",
    "        # self.net = nn.Sequential(list of layers)\n",
    "        # do not place ReLU at the end of the network!\n",
    "        # note: the first layer of model should have nin * bins input units\n",
    "        \n",
    "        # ====\n",
    "\n",
    "        self.create_mask()  # builds the initial self.m connectivity\n",
    "\n",
    "        \n",
    "    def create_mask(self):\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) The ordering of input units from 1 to m (self.ordering).\n",
    "        # 2) Assign the random number k from 1 to m − 1 to each hidden unit. \n",
    "        #    This number gives the maximum number of input units to which the unit can be connected.\n",
    "        # 3) Each hidden unit with number k is connected with the previous layer units \n",
    "        #   which has the number is less or equal than k.\n",
    "        # 4) Each output unit with number k is connected with the previous layer units \n",
    "        #    which has the number is less than k.\n",
    "        \n",
    "        # ====\n",
    "\n",
    "        # set the masks in all MaskedLinear layers\n",
    "        layers = [l for l in self.net.modules() if isinstance(l, MaskedLinear)]\n",
    "        for l, m in zip(layers, masks):\n",
    "            l.set_mask(m)\n",
    "\n",
    "    def visualize_masks(self):\n",
    "        prod = self.masks[0]\n",
    "        for idx, m in enumerate(self.masks):\n",
    "            plt.figure(figsize=(3, 3))\n",
    "            plt.title(f'layer: {idx}')\n",
    "            plt.imshow(m, vmin=0, vmax=1, cmap='gray')\n",
    "            plt.show()\n",
    "\n",
    "            if idx > 0:\n",
    "                prod=prod.dot(m)\n",
    "\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.title('prod')\n",
    "        plt.imshow(prod, vmin=0, vmax=1, cmap='gray')\n",
    "        plt.show()\n",
    "\n",
    "    def forward(self, x):\n",
    "        assert len(x.size()) == 2\n",
    "        assert x.shape[1] == self.nin\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply one hot encoding to x\n",
    "        # 2) apply the model\n",
    "        # 3) reshape and transpose the output to (batch_size, self.bins, self.nin)\n",
    "        \n",
    "        # ====\n",
    "\n",
    "    def sample(self, n, use_cuda=True):\n",
    "        # read carefully and understand the sampling process\n",
    "        xs = []\n",
    "        for _ in range(n):\n",
    "            x = torch.randint(0, self.bins, (1, self.nin))\n",
    "            if use_cuda:\n",
    "                x = x.cuda()\n",
    "            for it in range(self.nin):\n",
    "                probs = F.softmax(model(x)[0], dim=0).T\n",
    "                distr = torch.distributions.categorical.Categorical(probs)\n",
    "                x[0, it] = distr.sample()[it]\n",
    "            xs.append(x)\n",
    "        xs = torch.cat(xs)\n",
    "        return xs.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5WyKyynRGykw"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "HIDDEN_SIZES = \n",
    "# ====\n",
    "\n",
    "model = MADE(2, BINS, HIDDEN_SIZES)\n",
    "\n",
    "\n",
    "def test_model_output(model):\n",
    "    assert [10, BINS, 2] == list(model(torch.randint(0, BINS, (10, 2))).size())\n",
    "\n",
    "\n",
    "def test_create_mask(model):\n",
    "    prod = np.ones((1, BINS * 2))\n",
    "    for m in model.masks:\n",
    "        assert set(np.unique(m)).issubset((True, False))\n",
    "        prod = prod.dot(m)\n",
    "    assert np.allclose(prod, np.repeat(np.array([[0, BINS * np.prod(HIDDEN_SIZES)]]), BINS))\n",
    "\n",
    "test_create_mask(model)\n",
    "test_model_output(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zKam2n--SOPs"
   },
   "source": [
    "Now we will visualize the model masks. It should helps you to understand whether the model is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "svPeNKPbJsTP",
    "outputId": "7637d75a-e610-4349-d79c-e4f2678f606c"
   },
   "outputs": [],
   "source": [
    "model.visualize_masks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_NWRhm_zdGGK"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# you have to choose these parameters by yourself\n",
    "BATCH_SIZE = \n",
    "EPOCHS = \n",
    "LR = \n",
    "# ====\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_cuda=USE_CUDA)\n",
    "\n",
    "assert test_losses[-1] < 4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 738
    },
    "id": "uZRJq622LK4q",
    "outputId": "60bc0bb3-ad6c-48f0-ebf9-93531f32c456"
   },
   "outputs": [],
   "source": [
    "def get_distribution(model, use_cuda=True):\n",
    "    x = np.mgrid[0:model.bins, 0:model.bins].reshape(2, model.bins ** 2).T\n",
    "    # ====\n",
    "    # your code\n",
    "    # 1) take the model output for the grid x (shape: bins ** 2, bins, 2)\n",
    "    # 3) apply log_softmax to get log probs (shape: bins ** 2, bins, 2)\n",
    "    # 4) apply torch.gather to gather vaalues indexed by grid x (shape: bins ** 2, 2)\n",
    "    # 5) sum the log probs over dim=1 (shape: bins ** 2)\n",
    "    # 6) exponentiate it (shape: bins ** 2)\n",
    "    # 7) return an array BINS x BINS with probabilities of each pixel\n",
    "\n",
    "    \n",
    "    # ====\n",
    "\n",
    "\n",
    "distribution = get_distribution(model, USE_CUDA)\n",
    "assert distribution.shape == (BINS, BINS)\n",
    "\n",
    "plot_training_curves(train_losses, test_losses)\n",
    "plot_2d_distribution(image, distribution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 301
    },
    "id": "gCRYHJALfMMk",
    "outputId": "35e07e4d-01b3-489e-cb0f-f76882d9d47c"
   },
   "outputs": [],
   "source": [
    "# draw samples from model \n",
    "\n",
    "samples = model.sample(5000)\n",
    "plot_2d_data(train_data, samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tdgaPdSxQqXp"
   },
   "source": [
    "## Task 3: MADE on MNIST (3pt)\n",
    "\n",
    "\n",
    "You do not have to change this functions (except the path to the data file, download it from here: https://drive.google.com/file/d/1Ms-RBybrueI3_w2CRj7lM9mYjfvFRL6w/view?usp=sharing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "gHk4StQKLQ9n",
    "outputId": "3af7e8d9-55f9-4ea5-a901-7db758d42aa1"
   },
   "outputs": [],
   "source": [
    "# change the path to the file\n",
    "train_data, test_data = load_pickle(os.path.join('drive', 'My Drive', 'DGM', 'homework_supplementary', 'mnist.pkl'))\n",
    "visualize_mnist_images(train_data, 'MNIST samples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HFyUdMB5RQJm"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "HIDDEN_SIZES = \n",
    "# ====\n",
    "\n",
    "BINS = 2\n",
    "\n",
    "model = MADE(28 * 28, BINS, HIDDEN_SIZES)\n",
    "\n",
    "def test_model_output(model):\n",
    "    assert [10, BINS, 28 * 28] == list(model(torch.randint(0, 2, (10, 28 * 28))).size())\n",
    "\n",
    "\n",
    "test_model_output(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 957
    },
    "id": "S5oSStf8SSf8",
    "outputId": "7f739ad7-29f3-4519-f302-67ae1985f5aa"
   },
   "outputs": [],
   "source": [
    "# show on your masks and assure that they are correct\n",
    "model.visualize_masks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ueL7uJPOSveZ",
    "outputId": "0eec6f79-23b5-4a2b-d366-cd65c9c1c497"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# you have to choose these parameters by yourself\n",
    "BATCH_SIZE = \n",
    "EPOCHS = \n",
    "LR = \n",
    "# ====\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[-1] < 0.16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 404
    },
    "id": "Qvt5VHqRTI84",
    "outputId": "5a2a8d38-a5a0-46e6-8cd3-c22336ab76ac"
   },
   "outputs": [],
   "source": [
    "plot_training_curves(train_losses, test_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "FJDGq9OAWphh",
    "outputId": "aef32844-1150-4e3f-a302-196e1244dd32"
   },
   "outputs": [],
   "source": [
    "samples = model.sample(25)\n",
    "show_samples(samples, title='MNIST samples', nrow=5)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "hw1.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}