diffusion_model_book_2_2_score_based_model.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/shu65/diffusion-model-book/blob/main/diffusion_model_book_2_2_score_based_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ybGLrN9MeBRz"
      },
      "outputs": [],
      "source": [
        "from tqdm import tqdm_notebook as tqdm\n",
        "import torch\n",
        "\n",
        "device=\"cpu\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "zJNm-Rep08EY"
      },
      "outputs": [],
      "source": [
        "n_samples = int(1e6)\n",
        "sigma = 0.01\n",
        "\n",
        "dist0 = torch.distributions.MultivariateNormal(torch.tensor([-2, -2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))\n",
        "samples0 = dist0.sample(torch.Size([n_samples//2]))\n",
        "    \n",
        "dist1 = torch.distributions.MultivariateNormal(torch.tensor([2, 2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))\n",
        "samples1 = dist1.sample(torch.Size([n_samples//2]))\n",
        "samples = torch.vstack((samples0, samples1))\n",
        "\n",
        "mean = torch.mean(samples, dim=0)\n",
        "std = torch.std(samples, dim=0)\n",
        "\n",
        "normalized_samples = (samples - mean[None, :])/std[None, :]\n",
        "\n",
        "dataset = torch.utils.data.TensorDataset((normalized_samples))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 281
        },
        "id": "f8wwV52v4Cqr",
        "outputId": "dec54e06-aebc-4f21-e13a-07e2779eb5bf"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQoAAAEICAYAAACnA7rCAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/NK7nSAAAACXBIWXMAAAsTAAALEwEAmpwYAAAS0ElEQVR4nO3de7RU5X3G8e/DESFeEBBBBKKmIbbmVo0RXWprqzZqjeZmq81NkyybVNOkyyaxsc3tj8aupkmay2pqYxqjVs0yMcWKUYjWSxUrWryiCbUmgAgIghoQOef8+sd+DwzDDO/BPezZc3g+a5119t7znnnfrZxn3tn7PfNTRGBmtj2juj0AM6s/B4WZZTkozCzLQWFmWQ4KM8tyUJhZloPC2pL0BUlXdnscO0LScZKe6PY4RhoHRQ1JOlbS3ZLWSVoj6b8kvbXb49oRkp6StEHSC5LWpvP5qKSd+m8uIu6MiEOaxnHizuxzV+CgqBlJ44D/AL4JTASmAV8ENnZzXK/Q2yNib+BA4BLgM8Bl3R2SvRIOivp5HUBEXB0RAxGxISJuiYiHACT9hqRbJa2W9KykqySNH/rh9Ar6KUkPSfq1pMskTZF0U3p1nydpQmp7kKSQdJ6kpyUtl/SX7QYm6ag0M1gr6UFJxw/nhCJiXUTMBv4Y+KCkN6TnGyPpK5J+JWmFpO9IelV67HhJSyVdKGllGtu5DWM5VdJj6ZyWDY176OfS9hXAq4EbJL0o6dOSbpT08abzekjSO4dzLrusiPBXjb6AccBq4HLgFGBC0+OvBU4CxgD7AXcAX294/ClgPjCFYjayEngAOAwYC9wKfD61PQgI4GpgT+CNwCrgxPT4F4Ar0/a0NK5TKV5gTkr7+7U5j6eGnqfp+K+Aj6XtrwGzKWZOewM3AF9Ojx0P9ANfAkanftcP/fcAlgPHpe0JwOENP7e03TiAPwLubdh/czqP3bv9/77OX55R1ExEPA8cS/EL/C/AKkmzJU1Jjy+OiLkRsTEiVgFfBX636Wm+GRErImIZcCfFL8b/RMRLwPUUodHoixHx64h4GPhX4OwWQ3sfMCci5kTEYETMBRZQ/ALviKeBiZIEnAf8RUSsiYgXgL8Fzmpouwn4UkRsiog5wIvAIQ2PHSppXEQ8FxEPDLP/2cDrJM1M++8Hro2Il3fwPHYpDooaiohFEXFOREwH3gAcAHwdIL2NuCZNt58HrgQmNT3FiobtDS3292pqv6Rh+5epv2YHAmemtx1rJa2lCLSpO3RyxcxkDcVsaA/g/obn+2k6PmR1RPQ37K9vGPu7KULql5Jul3T0cDpPYXkt8L50YfVs4IodPIddjoOi5iLiceD7FIEBxatuAG+MiHEUr/Qq2c2Mhu1XU7zqN1sCXBER4xu+9oyIS4bbSbpzMw24C3iWIrRe3/B8+0REc4i1FBH3RcQZwGTgJ8AP2zVtcexy4L3ACcD6iLhnuOewq3JQ1Iyk30wX8Kan/RkUr3rzU5O9Kabg6yRNAz7VgW7/RtIekl4PnEvxitvsSuDtkt4mqU/S2HThcPowzmmcpNOAayiueTwcEYMUb62+JmlyajdN0tuG8Xy7S3qvpH0iYhPwPDDYpvkK4DWNB1IwDAL/gGcTw+KgqJ8XgFnAvZJ+TREQjwAXpse/CBwOrANuBH7cgT5vBxYDPwO+EhG3NDeIiCXAGcBnKS54LqEIqe39G7pB0gup7cUU11PObXj8M6nf+elt1Dy2XIPIeT/wVPq5j1LMEFr5MvDX6e1N4x2dH1BcvO2pBWXdonTl13ZBkg4C/g8Y3XQtYMST9AHgvIg4tttj6QWeUdguR9IewJ8Bl3Z7LL2idFBImiHptrT45VFJn2jRRpK+IWlxWtxyeNl+zV6JdA1kFcW1i3/r8nB6Rum3HpKmAlMj4gFJewP3A++IiMca2pwKfJzidtYs4B8jYlapjs2sMqVnFBGxfGixS1o0s4jiFlijM4AfRGE+MD4FjJn1gN06+WTp4thhwL1ND01j60U9S9Ox5S2e4zyKFXv00feWPRjXySGaWYMXeO7ZiNgv165jQSFpL+BHwCfTMuRXJCIuJV1kGqeJMUsndGiEZtZsXlz3y+G068hdD0mjKULiqohodV9/GVuv/puejplZD+jEXQ9RfMbAooj4aptms4EPpLsfRwHrImKbtx1mVk+deOtxDMUquYclLUzHPkvxNwNExHeAORR3PBZT/GHPuds+jZnVVemgiIi7yPxRUhT3YM8v25eZdYdXZppZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyOvUp3N+TtFLSI20eP17SOkkL09fnOtGvmVWjU3U9vg98i6KUfDt3RsRpHerPzCrUkRlFRNwBrOnEc5lZ/VR5jeJoSQ9KuknS6yvs18xK6mjt0e14ADgwIl5Mlc1/Asxs1bCx9uhY9qhoeGa2PZXMKCLi+Yh4MW3PAUZLmtSm7aURcUREHDGaMVUMz8wyKgkKSfun0oNIOjL1u7qKvs2svI689ZB0NXA8MEnSUuDzwGjYXFLwPcDHJPUDG4CzUvUwM+sBHQmKiDg78/i3KG6fmlkP8spMM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzrKo+4cqs/lS8bqqvr/g+ettfj9jUX3zv31TduGrAMwozy/KMwizNJEaNLT56cdS+EwEYnDJhm6Z9K54rHltVfEDb4MaXqhhh13lGYWZZnlHYLm/omsSoycXnPa/+nekArHrrtp/WuP/dewEw/q5iP55ZWXwf4dcsPKMws6yqao9K0jckLZb0kKTDO9GvmVWjqtqjp1AU/JkJzAL+KX0367qh26AD+xcXL1ceMwDAQ6d9c3ObF6O4LXpMXAjAuJ+PL342XdT0W49hGEbt0TOAH0RhPjBe0tRO9G1mO19VFzOnAUsa9pemY8ubG7qkoHVLFDWqiL7iIuYYbfn16Gt6jLS/q6jdxUyXFDSrn6pmFMuAGQ3709Mxs64bWpa92zNrAZh81wEAHNL3p1sabSpeU6fckxZnrSraDgwMVDPILqtqRjEb+EC6+3EUsC4itnnbYWb1VFXt0TnAqcBiYD1wbif6NeuESLOCwRXF4qlJtxfHJzy+zzZt+57Zegl37CIziqpqjwZwfif6MrPqeQm3WQwCMPjSxmL/6eJdsdIMo9HA0J+ZD80k0s+OdLW762Fm9eMZhdmQoZnFyy93eSD14xmFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW1amSgidLeiKVDLyoxePnSFolaWH6+kgn+jWzapT+4BpJfcC3gZMoCvvcJ2l2RDzW1PTaiLigbH9mVr1OzCiOBBZHxJMR8TJwDUUJQTMbIToRFO3KBTZ7d6pkfp2kGS0eB4qSgpIWSFqwiY0dGJ6ZlVXVxcwbgIMi4k3AXODydg1dUtCsfjoRFNlygRGxOiKGpgffBd7SgX7NrCKdCIr7gJmSDpa0O3AWRQnBzSRNbdg9HVjUgX7NrCKl73pERL+kC4CbgT7gexHxqKQvAQsiYjbw55JOB/qBNcA5Zfs1s+qoqPZXT+M0MWbphG4Pw2zEmhfX3R8RR+TaeWWmmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLKuqkoJjJF2bHr9X0kGd6NfMqlE6KBpKCp4CHAqcLenQpmYfBp6LiNcCXwP+rmy/ZladqkoKnsGWoj/XASdIUgf6NrMKVFVScHObiOgH1gH7tnoylxQ0q5/aXcx0SUGz+qmkpGBjG0m7AfsAqzvQt5lVoJKSgmn/g2n7PcCtUefKQ2a2lapKCl4GXCFpMUVJwbPK9mtm1SkdFAARMQeY03Tscw3bLwFndqIvM6te7S5mmln9OCjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCyrVFBImihprqRfpO8T2rQbkLQwfTV/8K6Z1VzZGcVFwM8iYibws7TfyoaI+O30dXrJPs2sYmWDorFU4OXAO0o+n5nVUNmgmBIRy9P2M8CUNu3GpjKB8yW9Y3tP6JKCZvWT/bh+SfOA/Vs8dHHjTkSEpHZFfQ6MiGWSXgPcKunhiPjfVg0j4lLgUoBxmugiQWY1kA2KiDix3WOSVkiaGhHLJU0FVrZ5jmXp+5OS/hM4DGgZFGZWP2XfejSWCvwg8O/NDSRNkDQmbU8CjgEeK9mvmVWobFBcApwk6RfAiWkfSUdI+m5q81vAAkkPArcBl0SEg8Ksh5QqKRgRq4ETWhxfAHwkbd8NvLFMP2bWXV6ZaWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmllXqr0dtO7Qlg9XXV3wf3fSfe2AAgMFN/cV+DFYyNLMd5RmFmWV5RrGTjBo7Zsv25EkA9E+duFWb3Z5ZW2ysXAXA4Pr1lYzNbEd5RmFmWZ5RdJh2Gw3AqP323XxszXHTAVhxXHFNIlQc3++eA4rvtxX7sfTpzT8T/Zt29lDNhq1sScEzJT0qaVDSEdtpd7KkJyQtltSumpiZ1VTZGcUjwLuAf27XQFIf8G3gJGApcJ+k2SP1A3aH7mwM7jd+87EVRxd3Mxb+4TcAGDfqVQAcPOojAExYVLQdla5VgGcUVi9lP1x3EYCk7TU7ElgcEU+mttdQlCIckUFhNhJVcY1iGrCkYX8pMKtdY0nnAecBjGWPnTuynWFw2+Jm6i+CdH0MpDYbiuObind+ivQz4cJoVk+lSgpGxDYFf8pySUGz+ilVUnCYlgEzGvanp2Nm1iOqeOtxHzBT0sEUAXEW8CcV9NsVkZZl9614bvOx/e/eC4BjuLBo01dMlCbfW7z16Fu+BoCBoaXcZjVT9vboOyUtBY4GbpR0czp+gKQ5ABHRD1wA3AwsAn4YEY+WG7aZVUlR4wto4zQxZmmbioU9YdSYsVu20+KrwSkTtm6zam1xfNXq4vuGDdUMziyZF9fdHxFt10AN8RJuM8vyEu6dZHDjS5u345mVAGj1mvRgMYsbSNczhq5rmNWVZxRmluUZRQWGlmN7Wbb1Ks8ozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpZVVUnBpyQ9LGmhpAVl+jSz6u30koINfi8ini3Zn5l1QRUlBc2sx1V1jSKAWyTdn0oGmlkPqaqk4LERsUzSZGCupMcj4o42/fV27VGzEaiKkoJExLL0faWk6ykqnLcMCtceNaufnf7WQ9KekvYe2gb+gOIiqJn1iJ1eUhCYAtwl6UHgv4EbI+KnZfo1s2qVvetxPXB9i+NPA6em7SeBN5fpx8y6yyszzSzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLKvvhun8v6XFJD0m6XtL4Nu1OlvSEpMWSLirTp5lVr+yMYi7whoh4E/Bz4K+aG0jqA74NnAIcCpwt6dCS/ZpZhUoFRUTcEhH9aXc+ML1FsyOBxRHxZES8DFwDnFGmXzOrVtlq5o0+BFzb4vg0YEnD/lJgVrsnaSwpCGycF9eNxGJBk4CRWNl9pJ4XjNxzO2Q4jTpSe1TSxUA/cNWOjLCVxpKCkhZExBFln7NufF69Z6Sem6QFw2lXuvaopHOA04ATIqJVrdBlwIyG/enpmJn1iLJ3PU4GPg2cHhHr2zS7D5gp6WBJuwNnAbPL9Gtm1Sp71+NbwN7AXEkLJX0Htq49mi52XgDcDCwCfhgRjw7z+S8tOb668nn1npF6bsM6L7V+t2BmtoVXZppZloPCzLJqHRTDXSLeiySdKelRSYOSev6220hdpi/pe5JWShpR63kkzZB0m6TH0r/DT2yvfa2DgmEsEe9hjwDvAu7o9kDKGuHL9L8PnNztQewE/cCFEXEocBRw/vb+n9U6KIa5RLwnRcSiiHii2+PokBG7TD8i7gDWdHscnRYRyyPigbT9AsUdyWnt2tc6KJp8CLip24Owllot02/7j87qRdJBwGHAve3adPJvPV6RqpeIV2k452bWTZL2An4EfDIinm/XrutB0YEl4rWVO7cRxMv0e5Ck0RQhcVVE/Hh7bWv91mOYS8St+7xMv8dIEnAZsCgivpprX+ugoM0S8ZFA0jslLQWOBm6UdHO3x/RKlVymX2uSrgbuAQ6RtFTSh7s9pg45Bng/8Pvpd2uhpFPbNfYSbjPLqvuMwsxqwEFhZlkOCjPLclCYWZaDwsyyHBRmluWgMLOs/wfC6EYofFebsQAAAABJRU5ErkJggg==\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "%matplotlib inline\n",
        "import numpy as np\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plot_samples = normalized_samples.cpu().numpy()\n",
        "\n",
        "plt.hist2d(plot_samples[:,0], plot_samples[:,1], range=((-2, 2), (-2, 2)), cmap='viridis', rasterized=False, bins=100, density=True)\n",
        "plt.gca().set_aspect('equal', adjustable='box')\n",
        "plt.xlim([-2, 2])\n",
        "plt.ylim([-2, 2])\n",
        "plt.title('Sample Density')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "x0gDBxSZdq8c"
      },
      "outputs": [],
      "source": [
        "sigma_begin = 0.001\n",
        "sigma_end = 1.0\n",
        "T = 200\n",
        "sigmas = torch.tensor(np.exp(np.linspace(np.log(sigma_begin), np.log(sigma_end), T))).float().to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "maXGG_Mdefn_"
      },
      "outputs": [],
      "source": [
        "def dsm_loss(score_model, samples, sigmas):\n",
        "  t = torch.randint(0, len(sigmas), (samples.shape[0],), device=sigmas.device)\n",
        "  used_sigmas = sigmas[t].view(samples.shape[0], *([1] * len(samples.shape[1:])))\n",
        "  noise = torch.randn_like(samples) * used_sigmas\n",
        "  perturbed_samples = samples + noise\n",
        "  target = - 1 / (used_sigmas ** 2) * noise\n",
        "  scores = score_model(perturbed_samples, used_sigmas)\n",
        "  target = target.view(target.shape[0], -1)\n",
        "  scores = scores.view(scores.shape[0], -1)\n",
        "  w = used_sigmas.squeeze(-1) ** 2\n",
        "  loss = ((scores - target) ** 2).sum(dim=-1) * w\n",
        "  return loss.mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "-Lwg4kMVM03Z"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class ScoreModel(nn.Module):\n",
        "  def __init__(self, n_channels=2):\n",
        "    super(ScoreModel, self).__init__()\n",
        "\n",
        "    self.model = nn.Sequential(\n",
        "        nn.Linear(n_channels, 2*n_channels),\n",
        "        nn.ELU(),\n",
        "        nn.Linear(2*n_channels, 16*n_channels),\n",
        "        nn.ELU(),\n",
        "        nn.Linear(16*n_channels, 2*n_channels),\n",
        "        nn.ELU(),\n",
        "        nn.Linear(2*n_channels, n_channels),\n",
        "    )\n",
        "\n",
        "  def forward(self, x, sigma):\n",
        "    y = self.model(x)\n",
        "    return y/sigma"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P9QWytgpa0pn",
        "outputId": "40654fc4-a598-4bc3-c450-9ee26158f83d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 steps loss:2.062662363052368\n",
            "1000 steps loss:1.9592573642730713\n",
            "2000 steps loss:1.7652524709701538\n",
            "3000 steps loss:1.6643431186676025\n",
            "4000 steps loss:1.6651298999786377\n",
            "5000 steps loss:1.811767578125\n",
            "6000 steps loss:1.7746050357818604\n",
            "7000 steps loss:1.7219853401184082\n",
            "8000 steps loss:1.6121973991394043\n",
            "9000 steps loss:1.5658706426620483\n",
            "10000 steps loss:1.5082696676254272\n",
            "11000 steps loss:1.6191147565841675\n",
            "12000 steps loss:1.369584083557129\n",
            "13000 steps loss:1.5210020542144775\n",
            "14000 steps loss:1.3806142807006836\n",
            "15000 steps loss:1.5240064859390259\n",
            "16000 steps loss:1.2849103212356567\n",
            "17000 steps loss:1.3859660625457764\n",
            "18000 steps loss:1.386408805847168\n",
            "19000 steps loss:1.4118430614471436\n",
            "20000 steps loss:1.270074725151062\n",
            "21000 steps loss:1.1704328060150146\n",
            "22000 steps loss:1.2353748083114624\n",
            "23000 steps loss:1.4386141300201416\n",
            "24000 steps loss:1.1812832355499268\n",
            "25000 steps loss:1.4304351806640625\n",
            "26000 steps loss:1.3862476348876953\n",
            "27000 steps loss:1.4564437866210938\n",
            "28000 steps loss:1.2314479351043701\n",
            "29000 steps loss:1.2191650867462158\n",
            "30000 steps loss:1.3120149374008179\n",
            "31000 steps loss:1.2689893245697021\n",
            "32000 steps loss:1.2594997882843018\n",
            "33000 steps loss:1.1706693172454834\n",
            "34000 steps loss:1.2825409173965454\n",
            "35000 steps loss:1.2294886112213135\n",
            "36000 steps loss:1.4062132835388184\n",
            "37000 steps loss:1.3233861923217773\n",
            "38000 steps loss:1.467130184173584\n",
            "39000 steps loss:1.2884862422943115\n",
            "40000 steps loss:1.3350117206573486\n",
            "41000 steps loss:1.327958106994629\n",
            "42000 steps loss:1.283301830291748\n",
            "43000 steps loss:1.328782558441162\n",
            "44000 steps loss:1.4744373559951782\n",
            "45000 steps loss:1.3854107856750488\n",
            "46000 steps loss:1.3049190044403076\n",
            "47000 steps loss:1.3279926776885986\n",
            "48000 steps loss:1.2053861618041992\n",
            "49000 steps loss:1.3855472803115845\n",
            "50000 steps loss:1.3562512397766113\n",
            "51000 steps loss:1.2336992025375366\n",
            "52000 steps loss:1.2602823972702026\n",
            "53000 steps loss:1.3706586360931396\n",
            "54000 steps loss:1.3223177194595337\n",
            "55000 steps loss:1.3518567085266113\n",
            "56000 steps loss:1.1787909269332886\n",
            "57000 steps loss:1.381345510482788\n",
            "58000 steps loss:1.2801939249038696\n",
            "59000 steps loss:1.4010896682739258\n",
            "60000 steps loss:1.3277764320373535\n",
            "61000 steps loss:1.363128900527954\n",
            "62000 steps loss:1.2995682954788208\n",
            "63000 steps loss:1.1924186944961548\n",
            "64000 steps loss:1.3905717134475708\n",
            "65000 steps loss:1.2209914922714233\n",
            "66000 steps loss:1.3428503274917603\n",
            "67000 steps loss:1.4009779691696167\n",
            "68000 steps loss:1.345718502998352\n",
            "69000 steps loss:1.1741654872894287\n",
            "70000 steps loss:1.3430674076080322\n",
            "71000 steps loss:1.4005886316299438\n",
            "72000 steps loss:1.3371143341064453\n",
            "73000 steps loss:1.39655339717865\n",
            "74000 steps loss:1.2602999210357666\n",
            "75000 steps loss:1.2962665557861328\n",
            "76000 steps loss:1.3343182802200317\n",
            "77000 steps loss:1.3735171556472778\n",
            "78000 steps loss:1.2625919580459595\n",
            "79000 steps loss:1.2656930685043335\n",
            "80000 steps loss:1.359081745147705\n",
            "81000 steps loss:1.2804900407791138\n",
            "82000 steps loss:1.3188785314559937\n",
            "83000 steps loss:1.3511507511138916\n",
            "84000 steps loss:1.4355740547180176\n",
            "85000 steps loss:1.3963481187820435\n",
            "86000 steps loss:1.2978711128234863\n",
            "87000 steps loss:1.2934417724609375\n",
            "88000 steps loss:1.3742218017578125\n",
            "89000 steps loss:1.3810871839523315\n",
            "90000 steps loss:1.3432632684707642\n",
            "91000 steps loss:1.1529685258865356\n",
            "92000 steps loss:1.3155038356781006\n",
            "93000 steps loss:1.2502846717834473\n",
            "94000 steps loss:1.3136907815933228\n",
            "95000 steps loss:1.2250237464904785\n",
            "96000 steps loss:1.3568994998931885\n",
            "97000 steps loss:1.2673358917236328\n",
            "98000 steps loss:1.3009989261627197\n",
            "99000 steps loss:1.3548177480697632\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "batch_size = 512\n",
        "n_steps = 100000\n",
        "\n",
        "dataloader = torch.utils.data.DataLoader(dataset, batch_size=512, shuffle=True, num_workers=0)\n",
        "dataloader_iter = iter(dataloader)\n",
        "\n",
        "score_model = ScoreModel().to(device)\n",
        "\n",
        "optimizer = torch.optim.Adam(score_model.parameters())\n",
        "lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.001, total_steps=n_steps)\n",
        "\n",
        "\n",
        "\n",
        "for i in range(n_steps):\n",
        "  try:\n",
        "    x = next(dataloader_iter)[0]\n",
        "  except StopIteration:\n",
        "    dataloader_iter = iter(dataloader)\n",
        "    x = next(dataloader_iter)[0]\n",
        "  x = x.to(device)\n",
        "\n",
        "  optimizer.zero_grad()\n",
        "  loss = dsm_loss(score_model, x, sigmas)\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "  lr_scheduler.step()\n",
        "  if (i % 1000) == 0:\n",
        "    print(f\"{i} steps loss:{loss}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "9P_sAS2Fpa-o"
      },
      "outputs": [],
      "source": [
        "def sbm_sample(n_samples, score_model, sigmas, alpha=0.1):\n",
        "    sigma_T = sigmas[-1]\n",
        "    x_0 = torch.randn(n_samples, 2)*sigma_T\n",
        "    x_tk = x_0\n",
        "    K = 200\n",
        "    for t in range(len(sigmas) -1, -1, -1):\n",
        "      sigma_t = sigmas[t]\n",
        "      alpha_t = alpha*(sigma_t**2)/(sigma_T**2)\n",
        "      print(f\"t:{t}, sigma_t:{sigma_t}, alpha_t:{alpha_t}\")\n",
        "      for k in range(K+1):\n",
        "        u_k = torch.randn(n_samples, 2)\n",
        "        if (k == K) and t == 0:\n",
        "          u_k[:, :] = 0.0\n",
        "        with torch.no_grad():\n",
        "          score = score_model(x_tk, sigma_t)\n",
        "          x_tk = x_tk + alpha_t * score + np.sqrt(2 * alpha_t) * u_k\n",
        "    return x_tk"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ixng3Ojiqp_J",
        "outputId": "0eb6e15c-74c4-4f5f-d5a5-e972a723bdcd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "t:199, sigma_t:1.0, alpha_t:0.10000000149011612\n",
            "t:198, sigma_t:0.965883195400238, alpha_t:0.09329303354024887\n",
            "t:197, sigma_t:0.9329304099082947, alpha_t:0.08703591674566269\n",
            "t:196, sigma_t:0.90110182762146, alpha_t:0.08119844645261765\n",
            "t:195, sigma_t:0.8703591227531433, alpha_t:0.07575250416994095\n",
            "t:194, sigma_t:0.8406652808189392, alpha_t:0.07067181169986725\n",
            "t:193, sigma_t:0.8119844794273376, alpha_t:0.06593187898397446\n",
            "t:192, sigma_t:0.7842822074890137, alpha_t:0.061509858816862106\n",
            "t:191, sigma_t:0.7575250267982483, alpha_t:0.057384420186281204\n",
            "t:190, sigma_t:0.731680691242218, alpha_t:0.053535666316747665\n",
            "t:189, sigma_t:0.7067181468009949, alpha_t:0.049945052713155746\n",
            "t:188, sigma_t:0.6826071739196777, alpha_t:0.0465952567756176\n",
            "t:187, sigma_t:0.6593188047409058, alpha_t:0.04347012937068939\n",
            "t:186, sigma_t:0.6368249654769897, alpha_t:0.040554605424404144\n",
            "t:185, sigma_t:0.6150985956192017, alpha_t:0.03783462941646576\n",
            "t:184, sigma_t:0.5941134095191956, alpha_t:0.035297077149152756\n",
            "t:183, sigma_t:0.5738441944122314, alpha_t:0.032929714769124985\n",
            "t:182, sigma_t:0.5542664527893066, alpha_t:0.030721131712198257\n",
            "t:181, sigma_t:0.5353566408157349, alpha_t:0.028660673648118973\n",
            "t:180, sigma_t:0.5170920491218567, alpha_t:0.026738420128822327\n",
            "t:179, sigma_t:0.4994505047798157, alpha_t:0.024945080280303955\n",
            "t:178, sigma_t:0.48241087794303894, alpha_t:0.02327202633023262\n",
            "t:177, sigma_t:0.4659525752067566, alpha_t:0.021711179986596107\n",
            "t:176, sigma_t:0.4500557780265808, alpha_t:0.020255019888281822\n",
            "t:175, sigma_t:0.4347013235092163, alpha_t:0.018896525725722313\n",
            "t:174, sigma_t:0.4198707044124603, alpha_t:0.017629140987992287\n",
            "t:173, sigma_t:0.40554606914520264, alpha_t:0.016446761786937714\n",
            "t:172, sigma_t:0.39171016216278076, alpha_t:0.015343685634434223\n",
            "t:171, sigma_t:0.3783462643623352, alpha_t:0.014314589090645313\n",
            "t:170, sigma_t:0.365438312292099, alpha_t:0.013354516588151455\n",
            "t:169, sigma_t:0.3529707193374634, alpha_t:0.012458832934498787\n",
            "t:168, sigma_t:0.34092849493026733, alpha_t:0.01162322424352169\n",
            "t:167, sigma_t:0.32929712533950806, alpha_t:0.01084365975111723\n",
            "t:166, sigma_t:0.31806257367134094, alpha_t:0.010116379708051682\n",
            "t:165, sigma_t:0.307211309671402, alpha_t:0.009437879547476768\n",
            "t:164, sigma_t:0.29673025012016296, alpha_t:0.008804883807897568\n",
            "t:163, sigma_t:0.2866067588329315, alpha_t:0.008214343339204788\n",
            "t:162, sigma_t:0.27682867646217346, alpha_t:0.007663411553949118\n",
            "t:161, sigma_t:0.2673841714859009, alpha_t:0.007149429526180029\n",
            "t:160, sigma_t:0.25826188921928406, alpha_t:0.006669920869171619\n",
            "t:159, sigma_t:0.24945081770420074, alpha_t:0.00622257124632597\n",
            "t:158, sigma_t:0.24094036221504211, alpha_t:0.0058052255772054195\n",
            "t:157, sigma_t:0.2327202409505844, alpha_t:0.005415871273726225\n",
            "t:156, sigma_t:0.2247805893421173, alpha_t:0.005052631255239248\n",
            "t:155, sigma_t:0.21711179614067078, alpha_t:0.004713753238320351\n",
            "t:154, sigma_t:0.20970463752746582, alpha_t:0.0043976036831736565\n",
            "t:153, sigma_t:0.20255018770694733, alpha_t:0.00410265801474452\n",
            "t:152, sigma_t:0.19563983380794525, alpha_t:0.0038274943362921476\n",
            "t:151, sigma_t:0.18896523118019104, alpha_t:0.0035707857459783554\n",
            "t:150, sigma_t:0.1825183480978012, alpha_t:0.003331294748932123\n",
            "t:149, sigma_t:0.17629140615463257, alpha_t:0.0031078660394996405\n",
            "t:148, sigma_t:0.17027691006660461, alpha_t:0.002899422775954008\n",
            "t:147, sigma_t:0.16446761786937714, alpha_t:0.0027049598284065723\n",
            "t:146, sigma_t:0.15885651111602783, alpha_t:0.0025235391221940517\n",
            "t:145, sigma_t:0.1534368395805359, alpha_t:0.002354286378249526\n",
            "t:144, sigma_t:0.1482020765542984, alpha_t:0.0021963855251669884\n",
            "t:143, sigma_t:0.14314588904380798, alpha_t:0.0020490745082497597\n",
            "t:142, sigma_t:0.13826221227645874, alpha_t:0.0019116438925266266\n",
            "t:141, sigma_t:0.1335451602935791, alpha_t:0.0017834309255704284\n",
            "t:140, sigma_t:0.12898902595043182, alpha_t:0.0016638169763609767\n",
            "t:139, sigma_t:0.12458833307027817, alpha_t:0.001552225323393941\n",
            "t:138, sigma_t:0.12033778429031372, alpha_t:0.0014481182442978024\n",
            "t:137, sigma_t:0.1162322461605072, alpha_t:0.0013509935233741999\n",
            "t:136, sigma_t:0.11226677894592285, alpha_t:0.0012603829381987453\n",
            "t:135, sigma_t:0.10843659937381744, alpha_t:0.0011758495820686221\n",
            "t:134, sigma_t:0.10473708808422089, alpha_t:0.0010969857685267925\n",
            "t:133, sigma_t:0.10116379708051682, alpha_t:0.0010234114015474916\n",
            "t:132, sigma_t:0.09771241247653961, alpha_t:0.0009547715890221298\n",
            "t:131, sigma_t:0.09437878429889679, alpha_t:0.000890735536813736\n",
            "t:130, sigma_t:0.09115888178348541, alpha_t:0.0008309941622428596\n",
            "t:129, sigma_t:0.08804883807897568, alpha_t:0.0007752598030492663\n",
            "t:128, sigma_t:0.08504489064216614, alpha_t:0.0007232633652165532\n",
            "t:127, sigma_t:0.08214343339204788, alpha_t:0.0006747543811798096\n",
            "t:126, sigma_t:0.07934096455574036, alpha_t:0.0006294988561421633\n",
            "t:125, sigma_t:0.07663410902023315, alpha_t:0.0005872786859981716\n",
            "t:124, sigma_t:0.07401960343122482, alpha_t:0.0005478902021422982\n",
            "t:123, sigma_t:0.07149428874254227, alpha_t:0.0005111432983539999\n",
            "t:122, sigma_t:0.06905513256788254, alpha_t:0.00047686113975942135\n",
            "t:121, sigma_t:0.06669919937849045, alpha_t:0.0004448783292900771\n",
            "t:120, sigma_t:0.06442363560199738, alpha_t:0.0004150404711253941\n",
            "t:119, sigma_t:0.06222570687532425, alpha_t:0.00038720385055057704\n",
            "t:118, sigma_t:0.06010276824235916, alpha_t:0.0003612342698033899\n",
            "t:117, sigma_t:0.05805225670337677, alpha_t:0.0003370064659975469\n",
            "t:116, sigma_t:0.056071698665618896, alpha_t:0.00031440352904610336\n",
            "t:115, sigma_t:0.054158713668584824, alpha_t:0.0002933166397269815\n",
            "t:114, sigma_t:0.052310992032289505, alpha_t:0.00027364399284124374\n",
            "t:113, sigma_t:0.05052630975842476, alpha_t:0.00025529079721309245\n",
            "t:112, sigma_t:0.04880251735448837, alpha_t:0.00023816856264602393\n",
            "t:111, sigma_t:0.047137532383203506, alpha_t:0.0002221947070211172\n",
            "t:110, sigma_t:0.04552935063838959, alpha_t:0.0002072921779472381\n",
            "t:109, sigma_t:0.043976034969091415, alpha_t:0.0001933891762746498\n",
            "t:108, sigma_t:0.04247571527957916, alpha_t:0.00018041864677798003\n",
            "t:107, sigma_t:0.041026581078767776, alpha_t:0.00016831803077366203\n",
            "t:106, sigma_t:0.03962688520550728, alpha_t:0.0001570290041854605\n",
            "t:105, sigma_t:0.038274943828582764, alpha_t:0.00014649714285042137\n",
            "t:104, sigma_t:0.036969125270843506, alpha_t:0.00013667163148056716\n",
            "t:103, sigma_t:0.03570786118507385, alpha_t:0.00012750514724757522\n",
            "t:102, sigma_t:0.034489624202251434, alpha_t:0.00011895342322532088\n",
            "t:101, sigma_t:0.03331294655799866, alpha_t:0.00011097524838987738\n",
            "t:100, sigma_t:0.032176416367292404, alpha_t:0.00010353217658121139\n",
            "t:99, sigma_t:0.03107866272330284, alpha_t:9.658833005232736e-05\n",
            "t:98, sigma_t:0.030018357560038567, alpha_t:9.011018119053915e-05\n",
            "t:97, sigma_t:0.028994228690862656, alpha_t:8.406653068959713e-05\n",
            "t:96, sigma_t:0.02800503931939602, alpha_t:7.842822378734127e-05\n",
            "t:95, sigma_t:0.02704959735274315, alpha_t:7.316807023016736e-05\n",
            "t:94, sigma_t:0.02612675167620182, alpha_t:6.826072058174759e-05\n",
            "t:93, sigma_t:0.025235392153263092, alpha_t:6.368249887600541e-05\n",
            "t:92, sigma_t:0.02437444217503071, alpha_t:5.9411344409454614e-05\n",
            "t:91, sigma_t:0.023542864248156548, alpha_t:5.542664439417422e-05\n",
            "t:90, sigma_t:0.022739658132195473, alpha_t:5.1709208491956815e-05\n",
            "t:89, sigma_t:0.021963853389024734, alpha_t:4.8241086915368214e-05\n",
            "t:88, sigma_t:0.021214518696069717, alpha_t:4.5005581341683865e-05\n",
            "t:87, sigma_t:0.020490746945142746, alpha_t:4.198707392788492e-05\n",
            "t:86, sigma_t:0.01979166828095913, alpha_t:3.9171016396721825e-05\n",
            "t:85, sigma_t:0.019116440787911415, alpha_t:3.654383181128651e-05\n",
            "t:84, sigma_t:0.01846424862742424, alpha_t:3.409284909139387e-05\n",
            "t:83, sigma_t:0.017834309488534927, alpha_t:3.1806262995814905e-05\n",
            "t:82, sigma_t:0.017225859686732292, alpha_t:2.9673023163923062e-05\n",
            "t:81, sigma_t:0.016638169065117836, alpha_t:2.7682868676492944e-05\n",
            "t:80, sigma_t:0.016070527955889702, alpha_t:2.582618617452681e-05\n",
            "t:79, sigma_t:0.015522253699600697, alpha_t:2.4094035325106233e-05\n",
            "t:78, sigma_t:0.014992684125900269, alpha_t:2.247805787192192e-05\n",
            "t:77, sigma_t:0.01448118221014738, alpha_t:2.0970464902347885e-05\n",
            "t:76, sigma_t:0.013987131416797638, alpha_t:1.9563985915738158e-05\n",
            "t:75, sigma_t:0.013509934768080711, alpha_t:1.8251834262628108e-05\n",
            "t:74, sigma_t:0.013049019500613213, alpha_t:1.7027690773829818e-05\n",
            "t:73, sigma_t:0.012603829614818096, alpha_t:1.588565282872878e-05\n",
            "t:72, sigma_t:0.012173827737569809, alpha_t:1.4820208889432251e-05\n",
            "t:71, sigma_t:0.01175849512219429, alpha_t:1.3826221220369916e-05\n",
            "t:70, sigma_t:0.01135733351111412, alpha_t:1.2898902241431642e-05\n",
            "t:69, sigma_t:0.010969857685267925, alpha_t:1.2033778148179408e-05\n",
            "t:68, sigma_t:0.010595601983368397, alpha_t:1.1226677997910883e-05\n",
            "t:67, sigma_t:0.010234113782644272, alpha_t:1.0473708243807778e-05\n",
            "t:66, sigma_t:0.00988495908677578, alpha_t:9.771241820999421e-06\n",
            "t:65, sigma_t:0.00954771600663662, alpha_t:9.115888133237604e-06\n",
            "t:64, sigma_t:0.009221978485584259, alpha_t:8.504488505423069e-06\n",
            "t:63, sigma_t:0.008907354436814785, alpha_t:7.934096174722072e-06\n",
            "t:62, sigma_t:0.008603464812040329, alpha_t:7.4019603744091e-06\n",
            "t:61, sigma_t:0.008309941738843918, alpha_t:6.905513146193698e-06\n",
            "t:60, sigma_t:0.008026433177292347, alpha_t:6.442362973757554e-06\n",
            "t:59, sigma_t:0.007752597332000732, alpha_t:6.010276592860464e-06\n",
            "t:58, sigma_t:0.007488104049116373, alpha_t:5.607170351140667e-06\n",
            "t:57, sigma_t:0.007232633884996176, alpha_t:5.231099294178421e-06\n",
            "t:56, sigma_t:0.006985879968851805, alpha_t:4.880251708527794e-06\n",
            "t:55, sigma_t:0.006747544277459383, alpha_t:4.5529354792961385e-06\n",
            "t:54, sigma_t:0.006517339497804642, alpha_t:4.247571268933825e-06\n",
            "t:53, sigma_t:0.00629498902708292, alpha_t:3.962688879255438e-06\n",
            "t:52, sigma_t:0.0060802241787314415, alpha_t:3.696912699524546e-06\n",
            "t:51, sigma_t:0.00587278651073575, alpha_t:3.4489621612010524e-06\n",
            "t:50, sigma_t:0.005672425962984562, alpha_t:3.2176417334994767e-06\n",
            "t:49, sigma_t:0.005478901322931051, alpha_t:3.00183614854177e-06\n",
            "t:48, sigma_t:0.00529197882860899, alpha_t:2.800504034894402e-06\n",
            "t:47, sigma_t:0.005111433565616608, alpha_t:2.612675189084257e-06\n",
            "t:46, sigma_t:0.004937048070132732, alpha_t:2.437444436509395e-06\n",
            "t:45, sigma_t:0.004768611863255501, alpha_t:2.2739659470971674e-06\n",
            "t:44, sigma_t:0.004605921916663647, alpha_t:2.121451643688488e-06\n",
            "t:43, sigma_t:0.004448782652616501, alpha_t:1.979166654564324e-06\n",
            "t:42, sigma_t:0.004297004546970129, alpha_t:1.8464248796590255e-06\n",
            "t:41, sigma_t:0.004150404594838619, alpha_t:1.7225859210157068e-06\n",
            "t:40, sigma_t:0.004008806310594082, alpha_t:1.6070528090494918e-06\n",
            "t:39, sigma_t:0.003872038796544075, alpha_t:1.4992684782555443e-06\n",
            "t:38, sigma_t:0.0037399372085928917, alpha_t:1.3987130387249636e-06\n",
            "t:37, sigma_t:0.00361234275624156, alpha_t:1.30490195715538e-06\n",
            "t:36, sigma_t:0.003489101305603981, alpha_t:1.2173827599326614e-06\n",
            "t:35, sigma_t:0.0033700643107295036, alpha_t:1.1357333278283477e-06\n",
            "t:34, sigma_t:0.0032550885807722807, alpha_t:1.0595601906970842e-06\n",
            "t:33, sigma_t:0.003144035581499338, alpha_t:9.884960263661924e-07\n",
            "t:32, sigma_t:0.0030367712024599314, alpha_t:9.221980121765228e-07\n",
            "t:31, sigma_t:0.0029331662226468325, alpha_t:8.603464607404021e-07\n",
            "t:30, sigma_t:0.0028330960776656866, alpha_t:8.026433420127432e-07\n",
            "t:29, sigma_t:0.0027364399284124374, alpha_t:7.488103506148036e-07\n",
            "t:28, sigma_t:0.0026430815923959017, alpha_t:6.985880531829025e-07\n",
            "t:27, sigma_t:0.002552908146753907, alpha_t:6.517340125355986e-07\n",
            "t:26, sigma_t:0.0024658110924065113, alpha_t:6.080224466131767e-07\n",
            "t:25, sigma_t:0.0023816856555640697, alpha_t:5.672426937053388e-07\n",
            "t:24, sigma_t:0.002300430089235306, alpha_t:5.291978482091508e-07\n",
            "t:23, sigma_t:0.0022219468373805285, alpha_t:4.937048174724623e-07\n",
            "t:22, sigma_t:0.0021461411379277706, alpha_t:4.605921901656984e-07\n",
            "t:21, sigma_t:0.00207292172126472, alpha_t:4.297004636555357e-07\n",
            "t:20, sigma_t:0.0020022003445774317, alpha_t:4.008806229194306e-07\n",
            "t:19, sigma_t:0.0019338917918503284, alpha_t:3.739937426416873e-07\n",
            "t:18, sigma_t:0.0018679136410355568, alpha_t:3.489101345621748e-07\n",
            "t:17, sigma_t:0.001804186380468309, alpha_t:3.255088643072668e-07\n",
            "t:16, sigma_t:0.0017426334088668227, alpha_t:3.0367712611223396e-07\n",
            "t:15, sigma_t:0.0016831803368404508, alpha_t:2.833096175436367e-07\n",
            "t:14, sigma_t:0.0016257556853815913, alpha_t:2.643081700171024e-07\n",
            "t:13, sigma_t:0.0015702900709584355, alpha_t:2.465810950980085e-07\n",
            "t:12, sigma_t:0.0015167169040068984, alpha_t:2.300430281820809e-07\n",
            "t:11, sigma_t:0.0014649713411927223, alpha_t:2.1461410426582006e-07\n",
            "t:10, sigma_t:0.0014149913331493735, alpha_t:2.002200574224844e-07\n",
            "t:9, sigma_t:0.001366716343909502, alpha_t:1.867913539399524e-07\n",
            "t:8, sigma_t:0.001320088398642838, alpha_t:1.742633344292699e-07\n",
            "t:7, sigma_t:0.001275051268748939, alpha_t:1.6257557433618786e-07\n",
            "t:6, sigma_t:0.001231550588272512, alpha_t:1.5167169920005108e-07\n",
            "t:5, sigma_t:0.0011895340867340565, alpha_t:1.4149912885841331e-07\n",
            "t:4, sigma_t:0.001148951007053256, alpha_t:1.3200885007336183e-07\n",
            "t:3, sigma_t:0.0011097524547949433, alpha_t:1.2315506126014952e-07\n",
            "t:2, sigma_t:0.0010718912817537785, alpha_t:1.1489509432749401e-07\n",
            "t:1, sigma_t:0.0010353218531236053, alpha_t:1.0718913756591064e-07\n",
            "t:0, sigma_t:0.0010000000474974513, alpha_t:1.0000001537946446e-07\n"
          ]
        }
      ],
      "source": [
        "samples_pred = sbm_sample(n_samples=100000, score_model=score_model, sigmas=sigmas)\n",
        "samples_pred = samples_pred.cpu().numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 281
        },
        "id": "dtTR7L35w4DK",
        "outputId": "0fee0f88-1920-430a-c8e1-b4b4388a1756"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQoAAAEICAYAAACnA7rCAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/NK7nSAAAACXBIWXMAAAsTAAALEwEAmpwYAAATYUlEQVR4nO3df5BdZX3H8feHEBJCQIxBjEkIKBGJP9GUkEGUCtqACIhYQysSkTJVMqjD2NLS0bHVEcYq6ITKL6lgkR8NolGCkCgpUA2SMEEJMRKpkA0xEEKACIIh3/5xnsTD5d59Njl3z713+bxm7uz58ex5nrO797PnPOfZfRQRmJn1Z6dON8DMup+DwsyyHBRmluWgMLMsB4WZZTkozCzLQdEhkr4t6Ytp+TBJK2uqNyTtX0ddO0rSLEl3dLod20PSPpI2SRrW6bYMBgdFPyT9TtIz6QdgXXpzj253PRFxe0QcMID2DOobSNIbJN0iaYOkjZKWSjp6sOobDJIWSfqjpKckPZnO4WxJIwaz3oh4KCJGR8TzpXacNph11slBkff+iBgNvA2YCvxLYwFJO9feqsHxQ2AB8CrglcCZwJMdbdGOmR0RuwPjgLOAmcB8Sepss3qXg2KAImINcBPwRth2CX+GpPuB+9O2YyQtS7+NfybpzVs/X9JBku5Ov+muBUaW9h0uqa+0PlHS9yQ9KukxSXMkHQhcBExPVzgbU9kRkv5d0kPpquciSbuWjvVZSWslPSzp1FbnJ2kssB9waUQ8l17/GxF3pP0vl/Sj1KbH0/KE0ucvkvTFdN6bJP1Q0iskXZV+s98lad9S+ZB0pqQHJK2X9BVJTX8eJb1e0oJ0pbNS0l8P8Hv2h4hYBBwLTAfel463U7rK+G36+l4naUzat29q2ynpa7pe0jmlthwsaUk6p3WSvtbweTtL+hJwGDAnfS3mSLpQ0lcbzmuepM8M5Fw6LiL8avECfgccmZYnAsuBf0vrQfHbdwywK3AQ8AgwDRgGnJI+fwSwC/Ag8BlgOHAi8Cfgi+lYhwN9aXkYcA9wPrAbRaC8I+2bBdzR0MbzgXmpHbtTXBV8Oe2bAayjCLfdgO+mdu/f5FxFEXg/Ao4H9m7Y/wrgg8CoVM9/A98v7V8ErAJeC7wMuA/4DXAksDNwJfCfpfIB3JravU8qe1rjeaZ2rwY+lo5zELAemNLie7Zo63Eatt8GnJeWPwUsBiak78/FwNVp376pbZem7+tbgGeBA9P+nwMnp+XRwCENn7dzs3YABwMPAzul9bHA041f5259dbwB3fyieKNvAjZSvNH/A9g17Qvg3aWy3ySFSGnbSuBdwDvTD4lK+35G86CYDjy69Qeu4Xjb3kBpXcAfgNeWtk0H/i8tXw6cW9r3OloERdo/AZgD/BbYkt5ck1uUfSvweGl9EXBOaf2rwE2l9fcDy0rrAcworX8S+EnjeQIfBm5vqPti4PMt2vWCN2hp+zUUV0sAK4AjSvvGUQT3zqU3/ITS/l8AM9PybcAXgLENx9/6eU2DolTve9LybGB+p3/GB/ryrUfe8RGxZ0RMiohPRsQzpX2rS8uTgLPSbcfGdGswEXh1eq2J9BOSPNiivonAgxGxeQBt24viN/zSUp0/TttJ9Zbb2KpOACKiLyJmR8Rr0/n8geJKAEmjJF0s6UFJT1K8YfZs6OVfV1p+psl6Y0dwY9te3aRZk4BpDV/Xv6XoR9ke44ENpWPeUDreCuB5YO9S+d+Xlp8utf3jFIH763Q7dcx2tOEK4CNp+SPAd7brDDrIQVFN+Y2/GvhSCpWtr1ERcTWwFhjf0Jm2T4tjrgb2adFB2vinvusp3oBvKNX5sig6X0n1ThxAnS+uKGI1cCGpT4aiU/AAYFpE7EFxlQTFVc2Oamzbw03KrAb+p+HrOjoiPjHQSiRNBN4O3F465lENxxwZRT9UvyLi/og4iaKz9zxgrqTdmhVtsu2/gOMkvQU4EPj+QM+h0xwU7XMp8PeSpqmwm6T3Sdqd4r52M3CmpOGSTqC4Z23mFxRv8HPTMUZKOjTtWwdMkLQLQERsSfWeL+mVAJLGS/qrVP46YJakKZJGAZ9v1fjUWfkFSfunzr6xwKkU9/JQ9Es8A2xMHX8tj7UdPpvqnUjRb3BtkzI/Al4n6eT0tRsu6S9S526/0lXQu4AfUHxd56ddFwFfkjQpldtL0nEDabCkj0jaK33tN6bNW5oUXQe8prwhIvqAuyiuJK5vuDrtag6KNomIJcDfUdzjP07RsTcr7XsOOCGtb6C47/5ei+M8T3E/vz/wENCXygP8lKJD9feS1qdt/5jqWpxuCRZS/OYnIm4CLkiftyp9bOU5ivvshRSPRO+l6MSblfZfQNG5t54iPH7cz7EG6gfAUmAZcCPwrcYCEfEU8F6KR5wPU9wSnEfRCdnKHElPUbxZLwCup+gP2fqG/jpFB/Atqdxiik7ogZgBLJe0KR1nZos3/NeBE9MTom+Utl8BvIkeuu2A1LlmVjdJQdFRuqrTbamTpHdS3IJMih568/mKwqwmkoZT3GJd1kshAW0IChWDg26VdJ+k5ZI+1aSMJH1D0ipJv5T0tqr1mvWS1KeykeJR7AUdbcwOaMfQ483AWRFxd+q4WyppQUTcVypzFDA5vaZRjDkY6D2hDUER8ZIaTh0RKygGj/WkylcUEbE2Iu5Oy09RPJMe31DsOODKKCymeP4+rmrdZlaPtv4xk4qx/AcBdzbsGs8LB9f0pW1rmxzjdOB0gGEMe/so9mhnE82s5CkeXx8Re+XKtS0oVPz59fXApyNih//iMCIuAS4B2ENjYpqOaFMLzazRwpjb72jdrdry1CP15l4PXBURzcYHrOGFo/AmpG1m1gPa8dRDFANlVkTE11oUmwd8ND39OAR4IiJedNthZt2pHbcehwInA7+StCxt+2fS3xVExEUUQ2ePphgd+DTFnwybWY+oHBRR/GOTfh91pcElZ1Sty8w6wyMzzSzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlnt+i/cl0t6RNK9LfYfLukJScvS63PtqNfM6tGueT2+DcwBruynzO0RcUyb6jOzGrXliiIibgM2tONYZtZ96uyjmC7pHkk3SXpDjfWaWUVtnXu0H3cDkyJik6Sjge9TzGz+IuW5R0cyqqbmmVl/armiiIgnI2JTWp4PDJc0tkXZSyJiakRMHc6IOppnZhm1BIWkV6WpB5F0cKr3sTrqNrPq2nLrIelq4HBgrKQ+4PPAcNg2peCJwCckbQaeAWam2cPMrAe0JSgi4qTM/jkUj0/NrAd5ZKaZZTkozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwGQjsVr5eol+6Zm9mA1TX3qCR9Q9IqSb+U9LZ21Gtm9ahr7tGjKCb8mQxMA76ZPpp1n3SL8eEVD2/bNEJ/AuDK108qNsSW2pvVSXXNPXoccGUUFgN7ShrXjrrNbPDVNaXgeGB1ab0vbVvbWNBTClq3OGn3h7Yt76pi1rormdSp5nRU13VmekpBs+5T1xXFGmBiaX1C2mbWfVL/w/ETDmm576WmriuKecBH09OPQ4AnIuJFtx1m1p3qmnt0PnA0sAp4GvhYO+o1G1Qv0auHZuqaezSAM9pRl5nVr+s6M82s+zgozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsq11TCs6QtDJNGXh2k/2zJD0qaVl6ndaOes2sHpX/Z6akYcCFwHsoJva5S9K8iLivoei1ETG7an1mVr92XFEcDKyKiAci4jngGoopBM1siGhHULSaLrDRB9NM5nMlTWyyHyimFJS0RNKSP/FsG5pnZlXV1Zn5Q2DfiHgzsAC4olVBTylo1n3aERTZ6QIj4rGI2Hp5cBnw9jbUa2Y1aUdQ3AVMlrSfpF2AmRRTCG4jaVxp9VhgRRvqNbOaVH7qERGbJc0GbgaGAZdHxHJJ/wosiYh5wJmSjgU2AxuAWVXrNbP6qJjtrzvtoTExTUd0uhlmQ9bCmLs0IqbmynlkppllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCyrrikFR0i6Nu2/U9K+7ajXzOpROShKUwoeBUwBTpI0paHYx4HHI2J/4HzgvKr1mll96ppS8Dj+POnPXOAISWpD3WZWg7qmFNxWJiI2A08Ar2h2ME8paNZ9uq4z01MKmnWfWqYULJeRtDPwMuCxNtRtZjWoZUrBtH5KWj4R+Gl088xDZvYCdU0p+C3gO5JWUUwpOLNqvWZWn8pBARAR84H5Dds+V1r+I/ChdtRlZvXrus5MM+s+Dgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCzLQWFmWQ4KM8uqFBSSxkhaIOn+9PHlLco9L2lZejX+410z63JVryjOBn4SEZOBn6T1Zp6JiLem17EV6zSzmlUNivJUgVcAx1c8npl1oapBsXdErE3Lvwf2blFuZJomcLGk4/s7oKcUNOs+2X/XL2kh8Komu84pr0RESGo1qc+kiFgj6TXATyX9KiJ+26xgRFwCXAKwh8Z4kiCzLpANiog4stU+SeskjYuItZLGAY+0OMaa9PEBSYuAg4CmQWFm3afqrUd5qsBTgB80FpD0ckkj0vJY4FDgvor1mlmNqgbFucB7JN0PHJnWkTRV0mWpzIHAEkn3ALcC50aEg8Ksh1SaUjAiHgOOaLJ9CXBaWv4Z8KYq9ZhZZ3lkppllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzLAeFmWU5KMwsy0FhZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZVX6D1e2ndSQy7GlM+0w205VpxT8kKTlkrZImtpPuRmSVkpaJanVbGJm1qWq3nrcC5wA3NaqgKRhwIXAUcAU4CRJUyrW25P2u3Mk+905ku+uvoPvrr6juMJovMow60JV/7nuCgBJ/RU7GFgVEQ+kstdQTEXo/8Rt1iPq6KMYD6wurfcB01oVlnQ6cDrASEYNbstq9uizowHY8LwnQLPeUmlKwYh40YQ/VXlKQbPuU2lKwQFaA0wsrU9I28ysR9Rx63EXMFnSfhQBMRP4mxrq7TqbDiumZp2tw4oNfjxqPaLq49EPSOoDpgM3Sro5bX+1pPkAEbEZmA3cDKwArouI5dWabWZ1UkT3dgPsoTExTS+asdDM2mRhzF0aES3HQG3lh/hmluWgMLMsB4WZZTkozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFmWg8LMshwUZpbloDCzrLqmFPydpF9JWiZpSZU6zax+Vf8L99YpBS8eQNm/jIj1Feszsw6oY0pBM+txdfVRBHCLpKVpykAz6yF1TSn4johYI+mVwAJJv46IpjOgD+W5R816VR1TChIRa9LHRyTdQDHDedOg8NyjZt1n0G89JO0mafety8B7KTpBzaxHDPqUgsDewB2S7gF+AdwYET+uUq+Z1avqU48bgBuabH8YODotPwC8pUo9ZtZZHplpZlkOCjPLclCYWZaDwsyyHBRmluWgMLMsB4WZZTkozCzLQWFmWQ4KM8tyUJhZloPCzLIcFGaW5aAwsywHhZllOSjMLMtBYWZZDgozy3JQmFlW1X+u+xVJv5b0S0k3SNqzRbkZklZKWiXp7Cp1mln9ql5RLADeGBFvBn4D/FNjAUnDgAuBo4ApwEmSplSs18xqVCkoIuKWiNicVhcDE5oUOxhYFREPRMRzwDXAcVXqNbN6VZ3NvOxU4Nom28cDq0vrfcC0VgcpTykIPLsw5g7FyYLGAkNxZvehel4wdM/tgIEUasvco5LOATYDV21PC5spTykoaUlETK16zG7j8+o9Q/XcJC0ZSLnKc49KmgUcAxwREc3mCl0DTCytT0jbzKxHVH3qMQP4B+DYiHi6RbG7gMmS9pO0CzATmFelXjOrV9WnHnOA3YEFkpZJugheOPdo6uycDdwMrACui4jlAzz+JRXb1618Xr1nqJ7bgM5Lze8WzMz+zCMzzSzLQWFmWV0dFAMdIt6LJH1I0nJJWyT1/GO3oTpMX9Llkh6RNKTG80iaKOlWSfeln8NP9Ve+q4OCAQwR72H3AicAt3W6IVUN8WH63wZmdLoRg2AzcFZETAEOAc7o73vW1UExwCHiPSkiVkTEyk63o02G7DD9iLgN2NDpdrRbRKyNiLvT8lMUTyTHtyrf1UHR4FTgpk43wppqNky/5Q+ddRdJ+wIHAXe2KtPOv/XYIXUPEa/TQM7NrJMkjQauBz4dEU+2KtfxoGjDEPGulTu3IcTD9HuQpOEUIXFVRHyvv7JdfesxwCHi1nkept9jJAn4FrAiIr6WK9/VQUGLIeJDgaQPSOoDpgM3Srq5023aURWH6Xc1SVcDPwcOkNQn6eOdblObHAqcDLw7vbeWSTq6VWEP4TazrG6/ojCzLuCgMLMsB4WZZTkozCzLQWFmWQ4KM8tyUJhZ1v8D2xy6Bci4DZMAAAAASUVORK5CYII=\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "%matplotlib inline\n",
        "import numpy as np\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.hist2d(samples_pred[:,0], samples_pred[:,1], range=((-2, 2), (-2, 2)), cmap='viridis', rasterized=False, bins=100, density=True)\n",
        "plt.gca().set_aspect('equal', adjustable='box')\n",
        "plt.xlim([-2, 2])\n",
        "plt.ylim([-2, 2])\n",
        "plt.title('Predicted Sample Density')\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyO0tdwuowfWNu0JbwudxOko",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}