week10_probmodels/homework.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9vC1UlrLY_w"
      },
      "source": [
        "# Диффузионные модели"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mryab/dl-hse-ami/blob/main/week10_probmodels/homework.ipynb)"
      ],
      "metadata": {
        "id": "VlsT1kWmV7n2"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68oNOYBFLY_x"
      },
      "source": [
        "#### Разработчик: Аким Котельников"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sz1i68n_LY_y"
      },
      "source": [
        "План задания: пишем свою диффузионную модель и тестируем её на SwissRoll-ах (50% баллов), далее запускаем её на MNIST (20%), затем пишем DDIM и проверям его на MNIST (30%). Цель задания: разобраться как работают диффузионные модели, разобраться в формулах. При этом проблем с обучением моделей почти быть не должно. Формат: дозаполнить \"todo\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EcHHLrnuLY_0"
      },
      "source": [
        "### Импорты и SwissRoll-ы"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget --quiet --show-progress \"https://raw.githubusercontent.com/mryab/dl-hse-ami/master/week10_probmodels/utils.py\""
      ],
      "metadata": {
        "id": "OCEuXFYOV-Fy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IcjQMkRkLY_1"
      },
      "outputs": [],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lEZ1oA4BLY_2"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import seaborn as sns\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from sklearn.datasets import make_swiss_roll\n",
        "from sklearn.utils import shuffle"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "SEED = 777\n",
        "\n",
        "np.random.seed(SEED)\n",
        "torch.manual_seed(SEED)\n",
        "torch.cuda.manual_seed_all(SEED)"
      ],
      "metadata": {
        "id": "gbi3yqCBWHT5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4NeWUCL0LY_3"
      },
      "outputs": [],
      "source": [
        "from sklearn.datasets import make_circles, make_swiss_roll\n",
        "\n",
        "\n",
        "def make_swiss_dataset(num_samples):\n",
        "    X0, _ = make_swiss_roll(num_samples // 2, noise=0.3, random_state=0)\n",
        "    X1, _ = make_swiss_roll(num_samples // 2, noise=0.3, random_state=0)\n",
        "    X0 = X0[:, [0, 2]]\n",
        "    X1 = X1[:, [0, 2]]\n",
        "    X1 = -X1\n",
        "    X, y = shuffle(\n",
        "        np.concatenate([X0, X1], axis=0),\n",
        "        np.concatenate([np.zeros(len(X0)), np.ones(len(X1))], axis=0),\n",
        "        random_state=0)\n",
        "    X = (X - X.mean(axis=0)) / X.std(axis=0)\n",
        "\n",
        "    return X, y\n",
        "\n",
        "X, y = make_swiss_dataset(2000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "auZANvzMLY_5"
      },
      "outputs": [],
      "source": [
        "sns.scatterplot(x=X[:, 0], y=X[:, 1], hue=y);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uncDovU4LY_6"
      },
      "source": [
        "## DDPM (0.5 баллов)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jWqy7MvWLY_8"
      },
      "source": [
        "В данной части вам предстоит написать собстевнную диффузионную модель (DDPM) и протестировать её на датасете выше. \n",
        "\n",
        "### Напоминание\n",
        "\n",
        "Напомним, что диффузионная модель состоит из прямого и обратного процесса. Прямой диффузионный процесс определяется как апостериорное распределение $q(x_{1:T}|x_0)$. Это распределение также является Марковской цепочкой, которая постепенно добавляет гауссовский шум к данному начальному объекту $x_0$. На каждом шаге шум добавляется с различной магнитудой, которая определяется расписанием дисперсий\n",
        " $\\{\\beta_1, ... \\beta_T\\}$. При правильном выборе расписания в пределе по числу шагов $T$ мы должны сойтись к шуму из $\\mathcal{N}(0, I)$. В качестве распределений $q$ берут нормальные распределения: \n",
        "$$\n",
        " q(x_t | x_{t - 1}) := \\mathcal{N}(x_t; \\sqrt{1 - \\beta_t}x_{t - 1}, \\beta_tI), \\ \\ \\ \\ \\ \\ \\ q(x_{1:T}|x_0) = \\prod_{t = 1}^T q(x_t | x_{t - 1})\n",
        "$$\n",
        "\n",
        "Теперь посмотрим со стороны обратного процесса. Обратный процесс расшумляет шум, пока не получится объект из изначального распределения. Таким образом, диффузионная модель - это вероятностная модель с латентными переменными вида $p_\\theta(x_0) := \\int p_\\theta(x_{0:T}) dx_{1:T}$, где промежуточные состояния $x_1, ..., x_T$ соответствуют зашумленным объектам, a $x_0$ - объект из распределения. Совместное распределение $p_\\theta(x_{0:T})$ называет обратным диффузионным процессом, который представляет собой Марковскую цепочку из гауссовских распределений $p_\\theta(x_{i-1}|x_{i})$: \n",
        "\n",
        "$$\n",
        "p(x_{0:T}) = p(x_0) \\prod_{t = 1}^Tp_{\\theta}(x_{t-1}|x_t) \\ \\ \\ \\ \\ \\ \\ \\ \\ p_\\theta(x_{T})=\\mathcal{N}(x_T | 0, I)\n",
        "$$\n",
        "$$\n",
        "  p_{\\theta}(x_{t - 1}|x_t):= \\mathcal{N}(x_{t - 1}; \\mu_{\\theta}(x_t, t), \\Sigma_{\\theta}(x_t, t))\n",
        "$$\n",
        "\n",
        "Вернемся к распределению $q(x_t | x_{t - 1})$.  Для того чтобы получить $x_t$, придется итеративно получать $x_1, ..., x_{t - 1}$. Однако это можно сделать более эффективно благодаря нормальным распределениям. Для этого обозначим $\\alpha_t := 1- \\beta_t$ и $\\bar{\\alpha}_t:= \\prod_{i = 1}^t\\alpha_i$, тогда \n",
        "$$\n",
        "q(x_t | x_0) = \\mathcal{N}(x_t;\\sqrt{\\bar{\\alpha}_t} x_0, (1-\\bar{\\alpha}_t)I) \\quad \\quad \\quad \\quad \\quad \\quad (1)\n",
        "$$\n",
        "\n",
        "Затем модель можно обучать, оптимизируя отдельные члены суммы вариационной нижней оценки $\\log p_{\\theta}(x_0)$:\n",
        "$$\n",
        "L_{VLB} = \\mathbb{E}_q [\\underbrace{D_\\text{KL}(q(\\mathbf{x}_T | \n",
        "\\mathbf{x}_0) \\parallel p_\\theta(\\mathbf{x}_T))}_{L_T} + + \\sum_{t=2}^T \n",
        "\\underbrace{D_\\text{KL}(q(\\mathbf{x}_{t-1} | \\mathbf{x}_t, \n",
        "\\mathbf{x}_0) \\parallel p_\\theta(\\mathbf{x}_{t-1} \n",
        "| \\mathbf{x}_t))}_{L_{t-1}} \\underbrace{- \\log p_\\theta(\\mathbf{x}_0 \n",
        "| \\mathbf{x}_1)}_{L_0}\n",
        "$$\n",
        "\n",
        "Для обучение нужно лишь выписать $q(\\mathbf{x}_{t-1} | \\mathbf{x}_t, \\mathbf{x}_0) = \\mathcal{N}(\\mathbf{x}_{t-1}; \\boldsymbol{\\mu}(\\mathbf{x}_t, \\mathbf{x}_0), \\tilde{\\beta}_t \\mathbf{I}) $: \n",
        "\n",
        "$$\n",
        "    \\boldsymbol{\\mu}(\\mathbf{x}_t, \\mathbf{x}_0) = \\frac{\\sqrt{\\alpha_t}(1 - \\bar{\\alpha}_{t-1})}{1 - \\bar{\\alpha}_t} \\mathbf{x}_t + \\frac{\\sqrt{\\bar{\\alpha}_{t-1}}\\beta_t}{1 - \\bar{\\alpha}_t} \\mathbf{x}_0 \\ \\ \\ \\ \\ \\ (2)\n",
        "$$\n",
        "$$\n",
        "    \\tilde{\\beta}_t = \\frac{1 - \\bar{\\alpha}_{t-1}}{1 - \\bar{\\alpha}_t} \\cdot \\beta_t  \\quad \\quad \\quad \\quad \\quad \\quad \\quad (3)\n",
        "$$\n",
        "\n",
        "\n",
        "За подробностями читайте статью [Denoising Diffusion Probabilistic Models (Ho et al. 2020)](https://arxiv.org/abs/2006.11239). \n",
        "\n",
        "Тем не менее в упомянутой статье было показано, что обучаясь на более простой лосс, получаются результаты лучше.  \n",
        "\n",
        "Итак, заметим, что\n",
        "$$\n",
        "x_t(x_0, \\epsilon) = \\sqrt{\\bar{\\alpha}_t} x_0 +  \\sqrt{1-\\bar{\\alpha}_t}\\epsilon, \\ \\ \\ \\epsilon \\sim \\mathcal{N}(0, I) \\quad \\quad \\quad \\quad \\quad \\quad \\quad (4)\n",
        "$$\n",
        "\n",
        "Тогда пускай наша модель с весами $\\theta$ будет предсказывать $\\epsilon$ из равенства выше, а именно обучаться, минимизируя данную функцию потерь: \n",
        "\n",
        "$$L^{simple}_t = \\mathbb{E}_{x_0, \\epsilon, t}\\bigg[ \\|\\epsilon - \\epsilon_{\\theta}(x_t, t)\\|^2\\bigg]$$\n",
        "\n",
        "Именно этот лосс вы должны будете использовать.\n",
        "\n",
        "\n",
        "Чтобы сэмплировать (обратный процесс), нам нужно получить $\\mu_{\\theta}(x_t, x_0)$ из $\\epsilon_{\\theta}(x_t, t)$. Для этого получите $\\hat{x}_0(\\epsilon_{\\theta}, x_t)$ из уравнения (4) и подставьте его в равенство (2)."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Переходим к заданию. Ниже будут представлены две вспомогательные функции, которые вам понадобятся."
      ],
      "metadata": {
        "id": "oPkC73BJKHse"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u9gHH9dPLY_-"
      },
      "outputs": [],
      "source": [
        "# some functions you will need\n",
        "\n",
        "import math\n",
        "\n",
        "\n",
        "# utility function. basicalle, returns arr[timesteps], where timesteps are indices. (look at class Diffusion)\n",
        "def _extract_into_tensor(arr, timesteps, broadcast_shape):\n",
        "    \"\"\"\n",
        "    Extract values from a 1-D torch tensor for a batch of indices.\n",
        "    :param arr: 1-D torch tensor.\n",
        "    :param timesteps: a tensor of indices into torche array to extract.\n",
        "    :param broadcast_shape: a larger shape of K dimensions witorch torche batch\n",
        "                            dimension equal to torche lengtorch of timesteps.\n",
        "    :return: a tensor of shape [batch_size, 1, ...] where torche shape has K dims.\n",
        "    \"\"\"\n",
        "    res = arr.to(device=timesteps.device)[timesteps].float()\n",
        "    while len(res.shape) < len(broadcast_shape):\n",
        "        res = res[..., None]\n",
        "    return res.expand(broadcast_shape)\n",
        "\n",
        "# out beta_t. we use linear scheduler\n",
        "def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):\n",
        "    \"\"\"\n",
        "    Get a pre-defined beta schedule for the given name.\n",
        "    The beta schedule library consists of beta schedules which remain similar\n",
        "    in the limit of num_diffusion_timesteps.\n",
        "    Beta schedules may be added, but should not be removed or changed once\n",
        "    they are committed to maintain backwards compatibility.\n",
        "    \"\"\"\n",
        "    scale = 1000 / num_diffusion_timesteps\n",
        "    beta_start = scale * 0.0001\n",
        "    beta_end = scale * 0.02 \n",
        "    if schedule_name == \"linear\":\n",
        "        # Linear schedule from Ho et al, extended to work for any number of\n",
        "        # diffusion steps.\n",
        "        return np.linspace(\n",
        "            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64\n",
        "        )\n",
        "    elif schedule_name == \"quad\":\n",
        "        betas = torch.linspace(beta_start ** 0.5, beta_end ** 0.5, num_diffusion_timesteps) ** 2\n",
        "        return betas.numpy()\n",
        "    elif schedule_name == \"sigmoid\":\n",
        "        betas = torch.linspace(-6, 6, num_diffusion_timesteps)\n",
        "        betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start\n",
        "        return betas.numpy()\n",
        "    else:\n",
        "        raise NotImplementedError(f\"unknown beta schedule: {schedule_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Класс Diffusion (0.3 балла из 0.5 за DDPM)\n",
        "Вам нужно дописать недостающие части ниже (помечены todo)."
      ],
      "metadata": {
        "id": "YFFqaztHJXTx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SXtIrJyGLZAA"
      },
      "outputs": [],
      "source": [
        "from tqdm import tqdm\n",
        "\n",
        "class Diffusion:\n",
        "    def __init__(\n",
        "        self,\n",
        "        *,\n",
        "        betas: np.array,\n",
        "        loss_type: str = \"mse\"\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Class that simulates Diffusion process. Does not store model or optimizer.\n",
        "        \"\"\"\n",
        "\n",
        "        betas = torch.from_numpy(betas).double()\n",
        "        self.betas = betas\n",
        "        assert len(betas.shape) == 1, \"betas must be 1-D\"\n",
        "        assert (betas > 0).all() and (betas <= 1).all()\n",
        "\n",
        "        self.num_timesteps = int(betas.shape[0])\n",
        "\n",
        "        alphas = # todo\n",
        "        self.alphas_cumprod = # todo\n",
        "        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), self.alphas_cumprod[:-1]], dim=0)\n",
        "        self.alphas_cumprod_next = torch.cat([self.alphas_cumprod[1:], torch.tensor([0.0]), ], dim=0)\n",
        "        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)\n",
        "\n",
        "        # calculations for diffusion q(x_t | x_{t-1})\n",
        "        self.sqrt_alphas_cumprod = self.alphas_cumprod.sqrt()\n",
        "        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)\n",
        "        self.log_one_minus_alphas_cumprod = torch.log(1.0 - self.alphas_cumprod)\n",
        "        self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)\n",
        "        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)\n",
        "\n",
        "\n",
        "        # calculations for posterior q(x_{t-1} | x_t, x_0)\n",
        "        self.posterior_variance = # todo, var from (3)\n",
        "\n",
        "        # log calculation clipped because posterior variance is 0.\n",
        "        self.posterior_log_variance_clipped = torch.log(\n",
        "            torch.cat([self.posterior_variance[1:2], self.posterior_variance[1:]], dim=0)\n",
        "        )\n",
        "        self.posterior_mean_coef1 = # todo, coef of xt from (2) \n",
        "        self.posterior_mean_coef2 = # todo, coef of x0 from (2) \n",
        "\n",
        "    def q_mean_variance(self, x0, t):\n",
        "        \"\"\"\n",
        "        Get mean and variance of distribution q(x_t | x_0). Use equation (1).\n",
        "        \"\"\"\n",
        "\n",
        "        mean = # todo\n",
        "        variance = # todo\n",
        "        log_variance = # todo\n",
        "        return mean, variance, log_variance\n",
        "\n",
        "    def q_posterior_mean_variance(self, x_start, x_t, t):\n",
        "        \"\"\"\n",
        "        Compute mean and variance of diffusion posterior q(x_{t-1} | x_t, x_0).\n",
        "        Use equation (2) and (3).\n",
        "        \"\"\"\n",
        "        assert x_start.shape == x_t.shape\n",
        "        posterior_mean = # todo\n",
        "        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)\n",
        "        posterior_log_variance_clipped = _extract_into_tensor(\n",
        "            self.posterior_log_variance_clipped, t, x_t.shape\n",
        "        )\n",
        "        return posterior_mean, posterior_variance, posterior_log_variance_clipped\n",
        "    \n",
        "    def q_sample(self, x_start, t, noise=None):\n",
        "        \"\"\"\n",
        "        Diffuse data for a given number of diffusion steps.\n",
        "        Sample from q(x_t | x_0).\n",
        "        \"\"\"\n",
        "        if noise is None:\n",
        "            noise = torch.randn_like(x_start)\n",
        "        return # todo\n",
        "\n",
        "    def p_mean_variance(self, model_output, x, t):\n",
        "        \"\"\"\n",
        "        Apply model to get p(x_{t-1} | x_t). Use Equation (2) and plug in \\hat{x}_0;\n",
        "        \"\"\"\n",
        "        model_variance = torch.cat([self.posterior_variance[1:2], self.betas[1:]], dim=0)\n",
        "        model_log_variance = torch.log(model_variance)\n",
        "        model_variance = _extract_into_tensor(model_variance, t, x.shape)\n",
        "        model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)\n",
        "\n",
        "        model_mean = # todo\n",
        "\n",
        "        return {\n",
        "            \"mean\": model_mean,\n",
        "            \"variance\": model_variance,\n",
        "            \"log_variance\": model_log_variance,\n",
        "            \"pred_xstart\": pred_xstart,\n",
        "        }\n",
        "\n",
        "    def _predict_xstart_from_eps(self, x_t, t, eps):\n",
        "        \"\"\"\n",
        "        Get \\hat{x0} from epsilon_{theta}. Use equation (4) to derive it.\n",
        "        \"\"\"\n",
        "        raise NotImplementedError\n",
        "        return # todo\n",
        "\n",
        "    def p_sample(self, model_output, x, t):\n",
        "        \"\"\"\n",
        "        Sample from p(x_{t-1} | x_t).\n",
        "        \"\"\"\n",
        "        out = # todo; get mean, variance of p(xt-1|xt)\n",
        "        noise = torch.randn_like(x)\n",
        "        nonzero_mask = # todo, no noise when t == 0\n",
        "\n",
        "        sample = out[\"mean\"] + nonzero_mask * torch.exp(0.5 * out[\"log_variance\"]) * noise\n",
        "        return {\"sample\": sample}\n",
        "    \n",
        "    def p_sample_loop(self, model, shape, y_dist):\n",
        "        \"\"\"\n",
        "        Samples a batch=shape[0] using diffusion model.\n",
        "        \"\"\"\n",
        "\n",
        "        x = torch.randn(*shape, device=model.device)\n",
        "        indices = list(range(self.num_timesteps))[::-1]\n",
        "\n",
        "        y = torch.multinomial(\n",
        "            y_dist,\n",
        "            num_samples=shape[0],\n",
        "            replacement=True\n",
        "        ).to(x.device)\n",
        "\n",
        "        for i in tqdm(indices):\n",
        "            t = torch.tensor([i] * shape[0], device=x.device)\n",
        "            with torch.no_grad():\n",
        "                model_output = model(x, t, y)\n",
        "                out = self.p_sample(\n",
        "                    model_output,\n",
        "                    x,\n",
        "                    t\n",
        "                )\n",
        "                x = out[\"sample\"]\n",
        "        return x, y\n",
        "    \n",
        "    def train_loss(self, model, x0, y):\n",
        "        \"\"\"\n",
        "        Calculates loss L^{simple}_t for the given model, x0.\n",
        "        \"\"\"\n",
        "        t = torch.randint(0, self.num_timesteps, size=(x0.size(0),), device=x0.device)\n",
        "        noise = # todo\n",
        "        x_t = # todo\n",
        "        model_output = # todo\n",
        "        loss = # todo\n",
        "        return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rQ4Mi-mlLZAB"
      },
      "outputs": [],
      "source": [
        "T = 100\n",
        "\n",
        "diffusion = Diffusion(\n",
        "    betas=get_named_beta_schedule('linear', T),\n",
        "    loss_type=\"mse\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GvZykOR5LZAB"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def show_noising(diffusion, X, y):\n",
        "    fig, axs = plt.subplots(1, 10, figsize=(40, 5))\n",
        "    for i,t in enumerate(range(0, diffusion.num_timesteps, 10)):\n",
        "        x = diffusion.q_sample(\n",
        "           x_start=torch.from_numpy(X),\n",
        "            t=torch.ones_like(torch.from_numpy(y)).long() * t,\n",
        "        )\n",
        "\n",
        "        sns.scatterplot(x=x[:,0], y=x[:,1], hue=y, ax=axs[i])\n",
        "        axs[i].set(title=t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKnZBT70LZAC"
      },
      "source": [
        "Давайте посмотрим, как зашумляются наши данные с увеличением $t$. Как думаете, достаточно ли $T = 100$ или надо увеличить? Как можно понять, сколько достаточно? (не оценивается)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5AfcaKEeLZAC"
      },
      "outputs": [],
      "source": [
        "show_noising(diffusion, X, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c22OaGgtLZAE"
      },
      "source": [
        "### Модель (0.1 балл из 0.5 за DDPM)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oGRw1RGKLZAE"
      },
      "source": [
        "Тут мы реализуем модель с весами $\\theta$, которая параметризует обратный процесс. Модель не должна быть сложной и большой. Достаточно только линейных слоев. Не забудьте учесть классы $y$ и шаги $t$. Модель предсказывает шум $\\epsilon: \\epsilon_{\\theta}(x_t, t, y)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SSYC7KmvLZAF"
      },
      "outputs": [],
      "source": [
        "class DiffModel(nn.Module):\n",
        "    def __init__(self, d_in, num_emb=2):\n",
        "        super().__init__()\n",
        "\n",
        "        self.hidden = 128\n",
        "        self.x_proj = # todo\n",
        "        self.t_proj = # todo\n",
        "        self.y_embed = # todo\n",
        "        self.layers = # todo\n",
        "\n",
        "    def forward(self, x, t, y):\n",
        "        '''\n",
        "        :x input, e.g. images\n",
        "        :t 1d torch.LongTensor of timesteps\n",
        "        :y 1d torch.LongTensor of class labels\n",
        "        '''\n",
        "        # todo\n",
        "        return self.layers(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gGID5Ek7LZAF"
      },
      "outputs": [],
      "source": [
        "model = DiffModel(d_in=2)\n",
        "model.device = torch.device('cpu') # достаточно cpu"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FnIuhHwTLZAG"
      },
      "source": [
        "### Обучение модели (0.1 балл из 0.5 за DDPM)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HfjHOndQLZAG"
      },
      "source": [
        "Наконец, обучим нашу модель. Ниже за вас написан класс `Trainer`, который хранит модель, диффузию и оптимайзер. Вам надо лишь дописать функцию `_run_step`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hl6OajqTLZAG"
      },
      "outputs": [],
      "source": [
        "from copy import deepcopy\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "\n",
        "class Trainer:\n",
        "    def __init__(\n",
        "        self,\n",
        "        diffusion: Diffusion,\n",
        "        model: nn.Module,\n",
        "        train_iter, # iterable that yields (x, y)\n",
        "        lr: float,\n",
        "        weight_decay: float,\n",
        "        steps: int,\n",
        "        device: torch.device = torch.device('cpu')\n",
        "    ):\n",
        "        self.diffusion = diffusion\n",
        "\n",
        "        self.train_iter = train_iter\n",
        "        self.steps = steps\n",
        "        self.init_lr = lr\n",
        "        self.model = model\n",
        "        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=weight_decay)\n",
        "        self.device = device\n",
        "        self.log_every = 100\n",
        "        self.print_every = 500\n",
        "\n",
        "    def _anneal_lr(self, step: int):\n",
        "        \"\"\"\n",
        "        Performs annealing of lr.\n",
        "        \"\"\"\n",
        "\n",
        "        frac_done = step / self.steps\n",
        "        lr = self.init_lr * (1 - frac_done)\n",
        "        for param_group in self.optimizer.param_groups:\n",
        "            param_group[\"lr\"] = lr\n",
        "\n",
        "    def _run_step(self, x: torch.FloatTensor, y: torch.LongTensor):\n",
        "        \"\"\"\n",
        "        A single training step.\n",
        "        Calculates loss for a single batch. \n",
        "        Then performs a single optimizer step and returns loss.\n",
        "        \"\"\"\n",
        "        # todo\n",
        "\n",
        "        return loss\n",
        "\n",
        "    def run_loop(self):\n",
        "        \"\"\"\n",
        "        Training loop.\n",
        "        \"\"\"\n",
        "        step = 0\n",
        "        curr_loss_gauss = 0.0\n",
        "\n",
        "        curr_count = 0\n",
        "        while step < self.steps:\n",
        "            x, y = next(self.train_iter)\n",
        "            batch_loss = self._run_step(x, y)\n",
        "\n",
        "            self._anneal_lr(step)\n",
        "\n",
        "            curr_count += len(x)\n",
        "            curr_loss_gauss += batch_loss.item() * len(x)\n",
        "\n",
        "            if (step + 1) % self.log_every == 0:\n",
        "                gloss = np.around(curr_loss_gauss / curr_count, 4)\n",
        "                if (step + 1) % self.print_every == 0:\n",
        "                    print(f'Step {(step + 1)}/{self.steps} Loss: {gloss}')\n",
        "                curr_count = 0\n",
        "                curr_loss_gauss = 0.0\n",
        "\n",
        "            step += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kOh210qELZAG"
      },
      "source": [
        "Теперь обернем наши данные в `DataLoader`. Для этого за вас написан `FastTensorDataLoader`. Также у нас идет обучение не по эпохам, а по итерациям, поэтому нам нужен \"бесконечный\" итератор."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_a6yQGSTLZAG"
      },
      "outputs": [],
      "source": [
        "from utils import FastTensorDataLoader\n",
        "\n",
        "\n",
        "def get_data_iter(X: np.array, y: np.array, batch_size: int = 512):\n",
        "    X = torch.from_numpy(X).float()\n",
        "    y = torch.from_numpy(y).long()\n",
        "    dataloader = FastTensorDataLoader(X, y, batch_size=batch_size, shuffle=True)\n",
        "    while True:\n",
        "        yield from dataloader\n",
        "    \n",
        "data_iter = get_data_iter(X, y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mr-YGEIuLZAH"
      },
      "outputs": [],
      "source": [
        "# you can change hyperparameters\n",
        "trainer = Trainer(\n",
        "    diffusion,\n",
        "    model,\n",
        "    train_iter=data_iter,\n",
        "    lr=0.01,\n",
        "    weight_decay=0.0,\n",
        "    steps=6000\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U0dtk6JaLZAH"
      },
      "outputs": [],
      "source": [
        "trainer.run_loop() # < 1min"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nf5sPAxILZAH"
      },
      "source": [
        "Теперь насэмплируем данные из нашей модели"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y552va6ELZAH"
      },
      "outputs": [],
      "source": [
        "from typing import List, Tuple\n",
        "\n",
        "\n",
        "def sample_synthetic(\n",
        "    diffusion: Diffusion,\n",
        "    model: nn.Module,\n",
        "    num_samples: int,\n",
        "    batch: int = 1000,\n",
        "    shape: Tuple = (2,),\n",
        "    y_dist: List[int] = [0.5, 0.5],\n",
        "    ddim: bool = False,\n",
        "):\n",
        "    sample_func = diffusion.p_sample_loop\n",
        "    if ddim: # for the last task\n",
        "        sample_func = diffusion.ddim_sample\n",
        "    res_x = []\n",
        "    res_y = []\n",
        "    num_sampled = 0\n",
        "    while num_sampled < num_samples:\n",
        "        x, y= diffusion.p_sample_loop(\n",
        "            model,\n",
        "            shape=(batch, *shape),\n",
        "            y_dist=torch.tensor(y_dist)\n",
        "        )\n",
        "        res_x.append(x.cpu())\n",
        "        res_y.append(y.cpu())\n",
        "        num_sampled += batch\n",
        "    \n",
        "    res_x = torch.cat(res_x, dim=0)\n",
        "    res_y = torch.cat(res_y, dim=0)\n",
        "    return res_x[:num_samples], res_y[:num_samples]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lCYZQbfILZAI"
      },
      "outputs": [],
      "source": [
        "Xs, ys = sample_synthetic(diffusion, model, 1000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMDKIkc5LZAI"
      },
      "outputs": [],
      "source": [
        "sns.scatterplot(x=Xs[:, 0], y=Xs[:, 1], hue=ys)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qSS_RMuuLZAI"
      },
      "source": [
        "Оцените на глаз, что получилось (должно быть неплохо). Как можно численно оценить качество насэмплированных данных (именно этих данных)?  Пункт не оценивается. "
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Покажете процесс расшумления аналагично тому, как мы это делали для прямого процесса"
      ],
      "metadata": {
        "id": "qoU_XRKCLBQJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# todo"
      ],
      "metadata": {
        "id": "LdA95T6bLKdf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G_KJR7vLLZAJ"
      },
      "source": [
        "## MNIST (0.2 балла)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bl6fIqYkLZAJ"
      },
      "source": [
        "Перейдём к обучению диффузионной модели на MNIST. За вас уже написана архитектура модели. В данной задаче надо лишь получить относительно *хороший* результат на датасете, используя класс `Diffusion`, а также посмотреть на разные расписание шума."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class InfiniteDataLoader(DataLoader):\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        # Initialize an iterator over the dataset.\n",
        "        self.dataset_iterator = super().__iter__()\n",
        "        \n",
        "    def __iter__(self):\n",
        "        return self\n",
        "\n",
        "    def __next__(self):\n",
        "        try:\n",
        "            batch = next(self.dataset_iterator)\n",
        "        except StopIteration:\n",
        "            # Dataset exhausted, use a new fresh iterator.\n",
        "            self.dataset_iterator = super().__iter__()\n",
        "            batch = next(self.dataset_iterator)\n",
        "        return batch"
      ],
      "metadata": {
        "id": "nMz3KEuxLC5K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rFiFelJwLZAK"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets.mnist import MNIST\n",
        "from torchvision.transforms import Compose, Lambda, Normalize, ToTensor\n",
        "\n",
        "transform = Compose([\n",
        "    ToTensor(),\n",
        "    Normalize((0.1307,), (0.3081,))\n",
        "])\n",
        "\n",
        "dataset = MNIST(\"./datasets\", download=True, train=True, transform=transform)\n",
        "loader = InfiniteDataLoader(dataset, 512, shuffle=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ej70iAHKLZAK"
      },
      "outputs": [],
      "source": [
        "def show_images(images, ys, title=\"\"):\n",
        "    \"\"\"Shows the provided images as sub-pictures in a square\"\"\"\n",
        "\n",
        "    # Converting images to CPU numpy arrays\n",
        "    if type(images) is torch.Tensor:\n",
        "        images = images.detach().cpu().numpy()\n",
        "        ys = ys.detach().cpu().numpy()\n",
        "\n",
        "    # Defining number of rows and columns\n",
        "    fig = plt.figure(figsize=(8, 8))\n",
        "    rows = int(len(images) ** (1 / 2))\n",
        "    cols = round(len(images) / rows)\n",
        "\n",
        "    # Populating figure with sub-plots\n",
        "    idx = 0\n",
        "    for r in range(rows):\n",
        "        for c in range(cols):\n",
        "            fig.add_subplot(rows, cols, idx + 1)\n",
        "\n",
        "            if idx < len(images):\n",
        "                plt.imshow(images[idx][0], cmap=\"gray\")\n",
        "                plt.title(f\"{int(ys[idx])}\")\n",
        "                plt.tick_params(bottom = False, labelbottom=False)\n",
        "                idx += 1\n",
        "    fig.suptitle(title, fontsize=30)\n",
        "\n",
        "    # Showing the figure\n",
        "    plt.show()\n",
        "\n",
        "def show_first_batch(loader):\n",
        "    for batch in loader:\n",
        "        show_images(batch[0][:16], batch[1][:16], \"Images in the first batch\")\n",
        "        break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JvHVOBxKLZAL"
      },
      "outputs": [],
      "source": [
        "show_first_batch(loader)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ymS7VXuALZAL"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(f\"Using device: {device}\\t\" + (f\"{torch.cuda.get_device_name(0)}\" if torch.cuda.is_available() else \"CPU\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F6UXgAnbLZAM"
      },
      "outputs": [],
      "source": [
        "from utils import MyUNet\n",
        "\n",
        "model_mnist = MyUNet().to(device)\n",
        "model_mnist.device = device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcWdlIb3LZAM"
      },
      "source": [
        "### Смотрим на расписание шума"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2PlNaBRVLZAM"
      },
      "source": [
        "Прежде чем начать, постройте на одной картинке графики $\\sqrt{\\bar{\\alpha_t}}$ (относительно $t$) для различных расписаний (linear, quad, sigmoid). Объясните, чем они отличаются. Чтобы лучше это понять визуализируйте зашумление картинок (функция `show_noising_mnist`). Советую выбрать $T = 1000$ (классический выбор, если домен -- картинки). Данный пункт оценивается."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rt5z549FLZAN"
      },
      "outputs": [],
      "source": [
        "def get_alpha_bar(betas):\n",
        "    # todo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wpRDr6f5LZAN"
      },
      "outputs": [],
      "source": [
        "ts = range(1000)\n",
        "\n",
        "# your plots here\n",
        "\n",
        "plt.legend([\"linear\", \"quad\", \"sigmoid\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2RTnAXNmLZAN"
      },
      "outputs": [],
      "source": [
        "# almost the same as in the task with SwissRolls\n",
        "\n",
        "def show_noising_mnist(diffusion, img):\n",
        "   # todo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0OklbFbrLZAN"
      },
      "outputs": [],
      "source": [
        "for sch in [\"linear\", \"quad\", \"sigmoid\"]:\n",
        "    diffusion_temp = Diffusion(\n",
        "       betas=get_named_beta_schedule(sch, 1000),\n",
        "        loss_type=\"mse\"\n",
        "    )\n",
        "    show_noising_mnist(diffusion_temp, dataset[0][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWJNsaX7LZAO"
      },
      "source": [
        "### Обучаем"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m1ywU_RGLZAO"
      },
      "outputs": [],
      "source": [
        "scheduler = # choose your pokemon\n",
        "\n",
        "diffusion_mnist = Diffusion(\n",
        "    betas=get_named_beta_schedule(scheduler, 1000),\n",
        "    loss_type=\"mse\"\n",
        ")\n",
        "\n",
        "# feel free to change hyperaparameters\n",
        "\n",
        "trainer_mnist = Trainer(\n",
        "    diffusion_mnist,\n",
        "    model_mnist,\n",
        "    train_iter=loader,\n",
        "    lr=0.001,\n",
        "    steps=1000,\n",
        "    weight_decay=0.0,\n",
        "    device=device\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-9iksepULZAP"
      },
      "outputs": [],
      "source": [
        "trainer_mnist.run_loop() # <15min on 2080Ti for the author's solution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ge0P-sUwLZAP"
      },
      "outputs": [],
      "source": [
        "Xs, ys = sample_synthetic(\n",
        "    diffusion_mnist,\n",
        "    model_mnist,\n",
        "    num_samples=16,\n",
        "    batch = 16,\n",
        "    shape=(1, 28, 28),\n",
        "    y_dist = [0.1 for _ in range(10)]\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YlmccDpsLZAW"
      },
      "outputs": [],
      "source": [
        "show_images(Xs, ys)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a33pasGULZAX"
      },
      "source": [
        "Оцените качество насэмплированных картинок."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yIbkLGeYLZAX"
      },
      "source": [
        "## DDIM (0.3 балла)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XDZNP8CuLZAX"
      },
      "source": [
        "**ВАЖНО:** В данном задании используется нотация из оригинальной статьи. Оно отличается от нотации выше. А именно, считайте, что ниже $\\alpha_t := \\sqrt{\\bar{\\alpha}_t}$. Также $\\sigma_t$ ниже тоже отличается (о том, что это, будет написано дальше)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUsCigzcLZAX"
      },
      "source": [
        "В данной задаче вам предстоит реализовать DDIM сэмплирования. Подробнее читайте тут: [Denoising Diffusion Implicit Models, Song et al., 2020](https://arxiv.org/abs/2010.02502). \n",
        "\n",
        "Давайте вкратце опишем, в чем смысл. Идея следущая: давайте изменим прямой диффузионный процесс так, чтобы используя предобученную DDPM, приближать новый обратный процесс за меньшее число шагов. В данном заданее мы не будем ускорять процесс сэмплирования, но реализуем DDIM.\n",
        "\n",
        "Чтобы не обучать новую модель, нам нужен прямой диффузионный процесс, у которого будет такая же функция потерь (MSE на шум), а обратный процесс все еще останется марковским. Оказалось, что существует целое семейство немарковских прямых процессов, удовлетворяющих этим требования. Это семейство имеет следующий вид (индуцируемое $\\sigma \\in \\mathbb{R}^T$):\n",
        "\n",
        "![ddim](https://i.imgur.com/lB2KaOR.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1I_Jmtm4LZAY"
      },
      "source": [
        "Распределения выше были выбраны так, чтобы $q(x_t | x_0)$ оставалось таким же, как раньше. Это можно проверить, используя формулу Байеса и свойства нормальных распределений."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cOXJf2NZLZAY"
      },
      "source": [
        "Тогда обратный пройесс можно переписать как $q_{\\sigma}(x_{t-1}|x_t, \\hat{x}_{0}(x_t))$ или же:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZoO5X8WLZAY"
      },
      "source": [
        "\\begin{equation*}\n",
        "    x_{t-1} = \\sqrt{\\alpha_{t-1}} \\underbrace{\\left(\\frac{x_t - \\sqrt{1 - \\alpha_t} \\epsilon_\\theta^{(t)}(x_t)}{\\sqrt{\\alpha_t}}\\right)}_{\\text{`` predicted } x_0 \\text{''}} + \\underbrace{\\sqrt{1 - \\alpha_{t-1} - \\sigma_t^2} \\cdot \\epsilon_\\theta^{(t)}(x_t)}_{\\text{``direction pointing to } x_t \\text{''}} + \\underbrace{\\sigma_t \\epsilon_t}_{\\text{random noise}}  \\quad \\quad \\quad \\quad \\quad \\quad (8)\n",
        "\\end{equation*}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7h1IBEYrLZAZ"
      },
      "source": [
        "Теперь давайте скажем, что $\\sigma_t(\\eta) = \\eta\\sqrt{(1 - \\alpha_{t - 1})(1 - \\alpha_t)}\\sqrt{1 - \\alpha_t / \\alpha_{t - 1}}$, тогда при $\\eta = 1$ равенство (8) превращается в DDPM сэмплирование. При $\\eta = 0$ у нас $\\sigma_t = 0$ и пропадает третья компонента стохастичности, что и называется DDIM сэмплирование. Таким образом, у нас есть детерминистичный процесс сэмплирования: при заданном начальном латенте мы всегда насэмлируем один и тот же $x_0$. Также можно достичь ускорения сэмплирования, выбирая лишь подмножество шагов $0 \\leq \\tau_1 \\leq ... \\leq \\tau_S \\leq T, \\ \\ \\ S < T$ и сэмплируя с помощью DDIM по ним.\n",
        "\n",
        "Поскольку процесс детерминистичный, то мы можем инвертировать равенство (8) и получить $x_t(x_{t - 1})$, то есть процесс зашумления (reversed DDIM).\n",
        "\n",
        "Все вышеописанное может быть полезно для осмысленных интерполяций в латентом пространстве (взяли две картинки, сделали обратный DDIM, получили два шума, как-то их проинтерполировали, расшумили DDIM, получили что-то среднее). \n",
        "\n",
        "\n",
        "Итак, ниже вам придется реализовать прямой и обратный DDIM. Однако мы не будем интерполировать латенты, а лишь зашумим-расшумим наши картинки. Важно отметить, что нам не нужно переучивать модель, чтобы пользоваться DDIM."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vsljIaqyLZAZ"
      },
      "outputs": [],
      "source": [
        "class DiffusionWithDDIM(Diffusion):\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        super(DiffusionWithDDIM, self).__init__(*args, **kwargs)\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def ddim_step(\n",
        "        self,\n",
        "        model_out: torch.FloatTensor,\n",
        "        x: torch.FloatTensor,\n",
        "        t: torch.LongTensor,\n",
        "        eta: float = 0.0\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Performs ddim step. Use equation (8).\n",
        "        \"\"\"\n",
        "\n",
        "        eps = model_out \n",
        "\n",
        "        alpha_bar = # todo\n",
        "        alpha_bar_prev = # todo\n",
        "        sigma = # todo\n",
        "\n",
        "        noise = torch.randn_like(x)\n",
        "        mean_pred = # todo\n",
        "        nonzero_mask = # todo, no noise when t == 0\n",
        "        sample = mean_pred + nonzero_mask * sigma * noise\n",
        "\n",
        "        return sample\n",
        "    \n",
        "    @torch.no_grad()\n",
        "    def ddim_sample(\n",
        "        self,\n",
        "        model: nn.Module,\n",
        "        shape: Tuple,\n",
        "        y_dist: torch.FloatTensor,\n",
        "        y: torch.LongTensor = None,\n",
        "        noise: torch.FloatTensor = None,\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Performs ddim sampling.\n",
        "        \"\"\"\n",
        "        if noise is None:\n",
        "            x = torch.randn(shape)\n",
        "        else:\n",
        "            x = noise\n",
        "\n",
        "        b = x.shape[0]\n",
        "        if y is None:\n",
        "            y = torch.multinomial(\n",
        "                y_dist,\n",
        "                num_samples=shape[0],\n",
        "                replacement=True\n",
        "            )\n",
        "\n",
        "        device = x.device\n",
        "        for t in reversed(range(self.num_timesteps)):\n",
        "            print(f'Sample timestep {t:4d}', end='\\r')\n",
        "            t_array = (torch.ones(b, device=device) * t).long()\n",
        "            model_out = model(x, t_array, y)\n",
        "            x = self.ddim_step(\n",
        "                model_out,\n",
        "                x,\n",
        "                t_array\n",
        "            )\n",
        "        print()\n",
        "        return x\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def ddim_reverse_step(\n",
        "        self,\n",
        "        model_out,\n",
        "        x,\n",
        "        t,\n",
        "        eta=0.0 # dummy\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Performs DDIM reverse step, i.e. xt from x_{t-1}. Use equation (8) to derive.\n",
        "        \"\"\"\n",
        "        assert eta == 0.0, \"Eta must be zero.\"\n",
        "\n",
        "        eps = model_out\n",
        "\n",
        "        alpha_bar_next = # todo\n",
        "        mean_pred = # todo\n",
        "\n",
        "        return mean_pred\n",
        "    \n",
        "    \n",
        "    @torch.no_grad()\n",
        "    def ddim_reverse_sample(\n",
        "        self,\n",
        "        model: nn.Module,\n",
        "        x: torch.FloatTensor,\n",
        "        y: torch.LongTensor,\n",
        "    ):\n",
        "        device = x.device\n",
        "        # todo\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GtwI1FluLZAa"
      },
      "outputs": [],
      "source": [
        "diffusion_ddim = DiffusionWithDDIM(\n",
        "    betas=get_named_beta_schedule(scheduler, 1000),\n",
        "    loss_type=\"mse\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vY5g_ErFLZAa"
      },
      "source": [
        "Теперь развернем картинки из прошлого задания и получим латенты $x_T$. И насэмплируем из этих латентов с помощью DDIM."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H4fzG7L0LZAb"
      },
      "outputs": [],
      "source": [
        "eps_reveresed = diffusion_ddim.ddim_reverse_sample(model_mnist, Xs.to(device), ys.to(device))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "28FfHeTnLZAb"
      },
      "outputs": [],
      "source": [
        "Xs_new, _ = diffusion_ddim.ddim_sample(\n",
        "    model_mnist,\n",
        "    shape = (1, 28, 28),\n",
        "    y_dist = torch.tensor([0.1] * 10),\n",
        "    noise = eps_reveresed.to(model_mnist.device),\n",
        "    y=ys.to(model_mnist.device)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fXU06GeaLZAb"
      },
      "outputs": [],
      "source": [
        "show_images(Xs, ys)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i_B0HQDKLZAb"
      },
      "outputs": [],
      "source": [
        "show_images(Xs_new, ys)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wva8ol7aLZAc"
      },
      "source": [
        "Оцените полученные сэмплы. Идеально ли они восстановились? Если нет, то почему? Пункт оценивается."
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "2a3lrJG0MeQd"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.9.7 ('base')",
      "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.7"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "a06af253165e97d0c1e75e8bf6d3252013856f30b8177e11b02d3fa36c37333d"
      }
    },
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}