Paper1_trainingequations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "83b0a6c0-7ce6-41d4-9186-55224860ca0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.datasets import make_swiss_roll\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4607304d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_batch(batch_size, device='cpu'):\n",
    "\n",
    "    data, _ = make_swiss_roll(batch_size)\n",
    "    data = data[:, [2,0]] / 10\n",
    "    data = data *np.array([1, -1])\n",
    "    return torch.from_numpy(data).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a13f8603",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "\n",
    "    def __init__(self, N=40, data_dim=2, hidden_dim=64):\n",
    "        super(MLP, self).__init__()\n",
    "\n",
    "        self.network_head = nn.Sequential(nn.Linear(data_dim, hidden_dim),\n",
    "                                          nn.ReLU(),\n",
    "                                          nn.Linear(hidden_dim, hidden_dim),\n",
    "                                          nn.ReLU(),)\n",
    "        \n",
    "        self.network_tail = nn.ModuleList([nn.Sequential(nn.Linear(hidden_dim, hidden_dim),\n",
    "                                                         nn.ReLU(),\n",
    "                                                         nn.Linear(hidden_dim, data_dim * 2),) for t in range(N)])\n",
    "        \n",
    "\n",
    "    def forward(self, x, t):\n",
    "\n",
    "        h = self.network_head(x) # [batch_size, hidden_dim]\n",
    "        tmp = self.network_tail[t](h) # [batch_size, data_dim * 2]\n",
    "        mu, h = torch.chunk(tmp, 2, dim=1)\n",
    "        var = torch.exp(h)\n",
    "        std = torch.sqrt(var)\n",
    "\n",
    "        return mu, std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "181a120b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 2])\n",
      "torch.Size([64, 2])\n"
     ]
    }
   ],
   "source": [
    "model = torch.load('model_paper1')\n",
    "t = 5\n",
    "x = torch.randn((64, 2))\n",
    "mu, std = model(x, t)\n",
    "\n",
    "print(mu.shape)\n",
    "print(std.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9e77eb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiffusionModel():\n",
    "    def __init__(self, T, model: nn.Module, device, dim=2):\n",
    "        self.betas = (torch.sigmoid(torch.linspace(-18, 10, T)) * (3e-1 - 1e-5) + 1e-5).to(device) # scaling data\n",
    "        self.alphas = 1-self.betas\n",
    "        self.alphas_bar = torch.cumprod(self.alphas, 0)\n",
    "\n",
    "        self.T = T\n",
    "        self.model = model\n",
    "        self.dim = dim\n",
    "\n",
    "\n",
    "    def forward_process(self, x0, t):\n",
    "\n",
    "        # param t: Number of diffusion steps\n",
    "        # param x0: data\n",
    "\n",
    "        assert t > 0, 't should be greater than 0'\n",
    "        assert self.T <= self.T, 'T should be lower or equal than {self.T}'\n",
    "\n",
    "        t = t - 1 # Because we start indexing at 0\n",
    "\n",
    "        mu = torch.sqrt(self.alphas_bar[t].cuda()) * x0\n",
    "        std = torch.sqrt(1 - self.alphas_bar[t])\n",
    "        epsilon = torch.randn_like(x0)\n",
    "        xt = mu + epsilon * std # data ~ N(mu, std)\n",
    "\n",
    "        std_q = torch.sqrt((1 - self.alphas_bar[t-1]) / (1 - self.alphas_bar[t]) * self.betas[t])\n",
    "        m1 = torch.sqrt(self.alphas_bar[t-1]) * self.betas[t] / (1 - self.alphas_bar[t])\n",
    "        m2 = torch.sqrt(self.alphas[t]) * (1 - self.alphas_bar[t-1]) / (1 - self.alphas_bar[t])\n",
    "        mu_q = m1 * x0 + m2 * xt\n",
    "\n",
    "        return mu_q, std_q, xt\n",
    "    \n",
    "    def reverse_process(self, xt, t):\n",
    "\n",
    "        # param t: Number of diffusion steps\n",
    "        # param x0: data\n",
    "\n",
    "        assert t > 0, 't should be greater than 0'\n",
    "        assert self.T <= self.T, 'T should be lower or equal than {self.T}'\n",
    "\n",
    "        t = t - 1 # Because we start indexing at 0\n",
    "\n",
    "        mu, std = self.model(xt, t)\n",
    "        epsilon = torch.randn_like(xt)\n",
    "\n",
    "        return mu, std, mu + epsilon * std # data ~ N(mu, std)\n",
    "    \n",
    "\n",
    "    def sample(self, batch_size, device):\n",
    "        \n",
    "        noise = torch.randn((batch_size, self.dim)).to(device)\n",
    "        x = noise\n",
    "\n",
    "        samples = [x]\n",
    "\n",
    "        for t in range(self.T, 0, -1):\n",
    "\n",
    "            if not t == 1:\n",
    "                _, _, x = self.reverse_process(x,t) # only reverse if t greater than 1\n",
    "\n",
    "            samples.append(x)\n",
    "\n",
    "        return samples[::-1]\n",
    "    \n",
    "    def get_loss(self, x0):\n",
    "\n",
    "        # param x0: batch [batch_size, self.dim]\n",
    "\n",
    "        t = torch.randint(2, 40+1, (1,))\n",
    "        mu_q, sigma_q, xt = self.forward_process(x0, t)\n",
    "        mu_p, sigma_p, xt_minus1 = self.reverse_process(xt.float(), t)\n",
    "\n",
    "        KL = torch.log(sigma_p) - torch.log(sigma_q) + (sigma_q**2 + (mu_q - mu_p)**2) / (2 * sigma_p**2)\n",
    "        K = - KL.mean() # should be maximized\n",
    "        loss = - K # should be minimized\n",
    "\n",
    "        return loss\n",
    "\n",
    "       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0e255576",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x0 = sample_batch(3_000)\n",
    "# mlp_model = torch.load('model_paper1')\n",
    "# model = DiffusionModel(40, mlp_model, device = 'cuda')\n",
    "# xT = model.forward_process(x0, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "86fb535e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(model, file_name, device):\n",
    "\n",
    "    font_size = 14\n",
    "    fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "    N = 5_000\n",
    "    x0 = sample_batch(N).to(device)\n",
    "    samples = model.sample(N, device=device)\n",
    "\n",
    "    data = [x0.cpu(), model.forward_process(x0, 20)[-1].cpu(), model.forward_process(x0, 40)[-1].cpu()]\n",
    "    for i in range(3):\n",
    "\n",
    "        plt.subplot(2, 3, 1+i)\n",
    "        plt.scatter(data[i][:, 0].data.numpy(), data[i][:, 1].data.numpy(), alpha=0.1, s=1)\n",
    "        plt.xlim([-2,2])\n",
    "        plt.ylim([-2,2])\n",
    "        plt.gca().set_aspect('equal')\n",
    "\n",
    "        if i == 0:\n",
    "            plt.ylabel(r'$q(\\mathbf{x}^{(0..T)})$', fontsize = font_size)\n",
    "            plt.title(r'$t=0$', fontsize = font_size)\n",
    "        elif i == 1:\n",
    "            plt.title(r'$t=\\frac{T}{2}$', fontsize = font_size)\n",
    "        elif i == 2:\n",
    "            plt.title(r'$t=T$', fontsize = font_size)\n",
    "\n",
    "    time_steps = [0, 20, 40]\n",
    "\n",
    "    for i in range(3):\n",
    "\n",
    "        plt.subplot(2, 3, 4+i)\n",
    "        plt.scatter(samples[time_steps[i]][:, 0].data.cpu().numpy(), samples[time_steps[i]][:, 1].data.cpu().numpy(),\n",
    "                    alpha=0.1, c='r', s=1)\n",
    "        plt.xlim([-2,2])\n",
    "        plt.ylim([-2,2])\n",
    "        plt.gca().set_aspect('equal')\n",
    "\n",
    "        if i == 0:\n",
    "            plt.ylabel(r'$p(\\mathbf{x}^{(0..T)})$', fontsize = font_size)\n",
    "\n",
    "\n",
    "    plt.savefig(file_name, bbox_inches = 'tight')\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "64e17daa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "250f99ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300000/300000 [1:08:53<00:00, 72.59it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[6.08629343079375,\n",
       " 1.1292982364010222,\n",
       " 5.820658242359925,\n",
       " 3.717994088199966,\n",
       " 1.2216253015541183,\n",
       " 1.106133287881369,\n",
       " 3.715602037454629,\n",
       " 1.2204815674954717,\n",
       " 1.1103958984481237,\n",
       " 1.1036039820559094,\n",
       " 1.0941890279809086,\n",
       " 1.4875844043947117,\n",
       " 1.0996790337136007,\n",
       " 6.0266806564422915,\n",
       " 3.3843460701752877,\n",
       " 1.1004319960024975,\n",
       " 1.085159769876299,\n",
       " 2.354788451850932,\n",
       " 6.148726904298986,\n",
       " 1.1030782514257012,\n",
       " 4.693061815538073,\n",
       " 1.1108697969778447,\n",
       " 5.939920778187928,\n",
       " 1.107880284524887,\n",
       " 1.703470035560936,\n",
       " 1.7016903904101275,\n",
       " 6.353447533332236,\n",
       " 3.711755484593728,\n",
       " 6.079645313353152,\n",
       " 5.2648823085685,\n",
       " 6.062612402068562,\n",
       " 1.1041196952593775,\n",
       " 4.063817403315206,\n",
       " 1.1127500554856233,\n",
       " 1.0827463991693216,\n",
       " 1.0999909045676604,\n",
       " 1.4850110073943907,\n",
       " 1.482771415794581,\n",
       " 6.0799197342241476,\n",
       " 6.024881823329591,\n",
       " 5.991397198151655,\n",
       " 3.7088566436017016,\n",
       " 2.0337345091046566,\n",
       " 1.4787199444098886,\n",
       " 5.2621452058144635,\n",
       " 1.098718674173503,\n",
       " 1.0921949004771865,\n",
       " 3.019581993326693,\n",
       " 6.206688781424399,\n",
       " 6.074558858773824,\n",
       " 1.1069569193910385,\n",
       " 1.0876915670035185,\n",
       " 5.816453427698998,\n",
       " 1.0913010168285526,\n",
       " 4.062625203958188,\n",
       " 5.9372072673644025,\n",
       " 1.699259271000044,\n",
       " 1.1245372815265802,\n",
       " 1.095423397344172,\n",
       " 4.690061359712714,\n",
       " 2.0313580430300857,\n",
       " 6.0769631441689285,\n",
       " 1.0840230342002732,\n",
       " 6.143396521994226,\n",
       " 1.696958646539964,\n",
       " 1.1230093726978512,\n",
       " 1.6943898601006764,\n",
       " 6.350460386328863,\n",
       " 1.2121480973043386,\n",
       " 1.113373186836655,\n",
       " 3.705888834465272,\n",
       " 1.2114652876623628,\n",
       " 5.8113110136120705,\n",
       " 5.988140866834908,\n",
       " 3.7051682149253784,\n",
       " 5.808502035853905,\n",
       " 6.348356697098078,\n",
       " 1.1178887515772111,\n",
       " 3.3794899355679364,\n",
       " 1.077847128503561,\n",
       " 1.0914902273069835,\n",
       " 3.376999544928838,\n",
       " 4.6871062459314725,\n",
       " 1.0920977220231225,\n",
       " 6.07105586019224,\n",
       " 6.060720349536903,\n",
       " 4.685952761889932,\n",
       " 4.380903700584187,\n",
       " 1.4755427372004517,\n",
       " 4.058202054430111,\n",
       " 6.141100257337769,\n",
       " 1.6904299621705345,\n",
       " 1.4735984966025226,\n",
       " 1.3247043571218529,\n",
       " 1.1213274680876766,\n",
       " 1.323331909738055,\n",
       " 1.1013455635782545,\n",
       " 3.015864933437198,\n",
       " 1.1185493000374276,\n",
       " 1.1038551419069995,\n",
       " 6.204898355502715,\n",
       " 1.1110057338006079,\n",
       " 1.6883749884116555,\n",
       " 1.2053943430856,\n",
       " 3.3733861479128073,\n",
       " 1.1092223494470475,\n",
       " 6.024069695644801,\n",
       " 2.6750305771378904,\n",
       " 4.377843006301319,\n",
       " 1.0949095799186652,\n",
       " 1.1074538204909752,\n",
       " 1.3201748377961444,\n",
       " 3.699329801339262,\n",
       " 6.344075939702395,\n",
       " 1.1037394227118116,\n",
       " 1.0710619147141642,\n",
       " 1.1019808019247852,\n",
       " 6.202852668833173,\n",
       " 1.070292812158859,\n",
       " 6.068505012177013,\n",
       " 1.0964602571913211,\n",
       " 1.6831696910037282,\n",
       " 2.6714646658302064,\n",
       " 6.343296560328097,\n",
       " 1.0837339604841951,\n",
       " 6.137572424009844,\n",
       " 5.668400907979862,\n",
       " 2.352565590580204,\n",
       " 5.258671165451565,\n",
       " 3.6978564164690533,\n",
       " 1.0944214237192782,\n",
       " 3.369883661464675,\n",
       " 1.079096490542952,\n",
       " 5.9353063109220345,\n",
       " 1.113415434503205,\n",
       " 2.66950616615517,\n",
       " 1.0974764889592121,\n",
       " 5.80255913929727,\n",
       " 4.683153210920529,\n",
       " 1.094086390868335,\n",
       " 2.665009991427108,\n",
       " 1.0920748486345742,\n",
       " 6.022549516716055,\n",
       " 2.0291999238302827,\n",
       " 5.982571428970891,\n",
       " 1.0872344849642315,\n",
       " 1.0863453149381044,\n",
       " 4.056396123986473,\n",
       " 1.0852579063677439,\n",
       " 1.088956663740011,\n",
       " 6.199563818138482,\n",
       " 1.1099281326482247,\n",
       " 6.064506419972177,\n",
       " 5.799981155538158,\n",
       " 3.0117418051672904,\n",
       " 6.198594554393239,\n",
       " 1.0885856122325537,\n",
       " 1.1003051951088825,\n",
       " 5.2561688561405555,\n",
       " 1.1077353321228363,\n",
       " 1.2025381174944758,\n",
       " 5.9329916738210065,\n",
       " 1.0903189434324543,\n",
       " 5.462996009950008,\n",
       " 6.339467956150835,\n",
       " 6.338336763470057,\n",
       " 1.0843956498537126,\n",
       " 1.3171312741640608,\n",
       " 2.026491057000878,\n",
       " 2.6623123541514717,\n",
       " 4.682868701461098,\n",
       " 5.98019411950717,\n",
       " 1.0935986563371354,\n",
       " 5.930788950254955,\n",
       " 2.0240780945157297,\n",
       " 5.460215327151223,\n",
       " 1.108211298735556,\n",
       " 2.021885033955895,\n",
       " 4.3752445339675745,\n",
       " 1.0857136103857115,\n",
       " 3.0083959423432027,\n",
       " 4.373255488187296,\n",
       " 6.059377005392917,\n",
       " 1.1997980421699737,\n",
       " 1.087948676960925,\n",
       " 5.9766189387796675,\n",
       " 4.980163716918088,\n",
       " 1.1012276492416528,\n",
       " 5.795563748973446,\n",
       " 1.0722491076740872,\n",
       " 1.0832222782567076,\n",
       " 1.079438942791523,\n",
       " 1.1968683464269534,\n",
       " 1.0822765813279425,\n",
       " 5.793155497993589,\n",
       " 1.0764155190031883,\n",
       " 6.061784707747505,\n",
       " 1.6809114361242132,\n",
       " 4.052227786508057,\n",
       " 1.0897429208955636,\n",
       " 4.051857866928225,\n",
       " 3.693966115237197,\n",
       " 2.6593170506503077,\n",
       " 3.692651362543424,\n",
       " 5.455910222410084,\n",
       " 1.0735561596459429,\n",
       " 4.679159729830252,\n",
       " 1.1923846036702486,\n",
       " 1.4668839990961564,\n",
       " 1.081342579384094,\n",
       " 1.1900306267340424,\n",
       " 2.6555691460196824,\n",
       " 4.047022684839785,\n",
       " 2.3495736172958805,\n",
       " 3.366266121178914,\n",
       " 6.072598076474144,\n",
       " 6.05832006504669,\n",
       " 1.0835648124617079,\n",
       " 1.0993215751346697,\n",
       " 4.67771905036349,\n",
       " 4.367515783226363,\n",
       " 1.0863265434911173,\n",
       " 4.365254668141998,\n",
       " 1.0822139917767712,\n",
       " 3.362375389787512,\n",
       " 6.196358938113586,\n",
       " 1.0765310899228562,\n",
       " 6.0526489747901,\n",
       " 4.044146457019678,\n",
       " 1.0698388236691498,\n",
       " 6.194214300109616,\n",
       " 1.0972941548990907,\n",
       " 1.0730086665255176,\n",
       " 1.06728050212958,\n",
       " 4.9777541454779275,\n",
       " 5.787546953494461,\n",
       " 2.0179772136165304,\n",
       " 5.784906035135012,\n",
       " 1.065330429607929,\n",
       " 2.3481644818579017,\n",
       " 2.3459836589399177,\n",
       " 2.6513581560223907,\n",
       " 4.975984315154012,\n",
       " 6.022073053693911,\n",
       " 4.973914395738408,\n",
       " 1.070725916179868,\n",
       " 1.0750268163243504,\n",
       " 1.1088694758186952,\n",
       " 2.016254672743953,\n",
       " 6.057993608876752,\n",
       " 1.0975401935231177,\n",
       " 6.019720424002031,\n",
       " 2.34188951387283,\n",
       " 5.666464310184451,\n",
       " 1.0673151084336323,\n",
       " 3.359501478224776,\n",
       " 5.253162624554901,\n",
       " 5.663907134551783,\n",
       " 1.0881004019085425,\n",
       " 5.452369922662289,\n",
       " 1.0869174368766357,\n",
       " 6.0159411791958135,\n",
       " 5.778839319255899,\n",
       " 1.4642184537692633,\n",
       " 3.6876568494613067,\n",
       " 5.6602369131941455,\n",
       " 1.0679923810679375,\n",
       " 1.085890687493985,\n",
       " 1.0615627867092667,\n",
       " 1.0930444951119465,\n",
       " 1.078260799857152,\n",
       " 1.05785741567543,\n",
       " 6.133041781544915,\n",
       " 6.1921124167067445,\n",
       " 2.337939233614294,\n",
       " 5.929518957961721,\n",
       " 3.356044649301445,\n",
       " 3.3547880326871353,\n",
       " 6.339212589694473,\n",
       " 5.97128584807642,\n",
       " 1.064151568364806,\n",
       " 6.068338170230274,\n",
       " 6.189435587051483,\n",
       " 1.1846535967535392,\n",
       " 6.33640102011944,\n",
       " 2.6487801381836427,\n",
       " 1.6781881533787577,\n",
       " 4.971606530374963,\n",
       " 6.013035279402767,\n",
       " 6.053920386261101,\n",
       " 1.1834612639942663,\n",
       " 1.3117242796717574,\n",
       " 1.104315049670443,\n",
       " 1.1810937693163603,\n",
       " 1.091892177304449,\n",
       " 1.0518554815749104,\n",
       " 1.6754683358405005,\n",
       " 3.3497480962308455,\n",
       " 1.0840347302347766,\n",
       " 6.050566864583478,\n",
       " 1.6730569722409665,\n",
       " 6.064052794047762,\n",
       " 6.129339673347064,\n",
       " 4.96891208641926,\n",
       " 4.9668887416401235,\n",
       " 1.0661092300252755,\n",
       " 5.924350703592322,\n",
       " 6.047010361807257,\n",
       " 6.126586302101976,\n",
       " 5.966706177204294,\n",
       " 6.1873107362286754,\n",
       " 1.0784780483459457,\n",
       " 1.0750454446138356,\n",
       " 2.646649684225858,\n",
       " 1.1751853988142307,\n",
       " 6.047454295537564,\n",
       " 6.0432253577363975,\n",
       " 6.184575446797508,\n",
       " 3.347451030597119,\n",
       " 6.0103862610189465,\n",
       " 3.3453367915313477,\n",
       " 6.007741782086695,\n",
       " 1.061426946595967,\n",
       " 6.181891784128904,\n",
       " 6.060032532247684,\n",
       " 5.775605302719185,\n",
       " 1.4611047255339422,\n",
       " 3.0016046061846096,\n",
       " 1.459325081775924,\n",
       " 1.058100317506714,\n",
       " 5.963190786218161,\n",
       " 2.6430022748654336,\n",
       " 1.4563180323317493,\n",
       " 1.1000909275900113,\n",
       " 4.963132518542109,\n",
       " 6.043132935615051,\n",
       " 3.3411035335509895,\n",
       " 1.0777898853840502,\n",
       " 5.4488382194934735,\n",
       " 1.080696291647455,\n",
       " 4.960880362723526,\n",
       " 1.6666392245522166,\n",
       " 1.0954658748788992,\n",
       " 5.922873147722523,\n",
       " 5.919809778670799,\n",
       " 1.0779192565062472,\n",
       " 5.654944519622212,\n",
       " 5.770680572582391,\n",
       " 1.3088204393733571,\n",
       " 4.678091052383391,\n",
       " 5.9163343121274306,\n",
       " 1.1690614701884365,\n",
       " 1.082444601599717,\n",
       " 4.0384245336321385,\n",
       " 5.767913688059977,\n",
       " 1.0446331982870694,\n",
       " 6.178351132674297,\n",
       " 6.039486886857548,\n",
       " 1.077682645853351,\n",
       " 6.056102548895116,\n",
       " 1.077198842042767,\n",
       " 1.0572839964395513,\n",
       " 6.33286637166274,\n",
       " 5.911577727932481,\n",
       " 1.042517798515437,\n",
       " 1.0695082377654594,\n",
       " 6.175238893282753,\n",
       " 6.329678436751011,\n",
       " 1.304870153297922,\n",
       " 1.0912557590563516,\n",
       " 1.0629752172246054,\n",
       " 5.9061932233249514,\n",
       " 2.9957213776842675,\n",
       " 5.957567769449729,\n",
       " 3.681747489641977,\n",
       " 6.051692548144647,\n",
       " 6.037491577627009,\n",
       " 6.326804323106963,\n",
       " 1.051888801592143,\n",
       " 6.172253180649537,\n",
       " 6.003852224642959,\n",
       " 1.0740316008382027,\n",
       " 1.0665317006172623,\n",
       " 1.6625072976750512,\n",
       " 2.9914868640643157,\n",
       " 1.0665946999877918,\n",
       " 1.301687264017208,\n",
       " 1.0674981491144035,\n",
       " 6.033199109387153,\n",
       " 1.3015764245074561,\n",
       " 4.957099399354653,\n",
       " 2.012604035651099,\n",
       " 5.764331529680964,\n",
       " 5.65170455679104,\n",
       " 3.678904253876659,\n",
       " 4.677158735479321,\n",
       " 1.0925389489069761,\n",
       " 1.6583480240832291,\n",
       " 1.0620322903149149,\n",
       " 1.057373003776495,\n",
       " 6.037308809595011,\n",
       " 5.999450778765355,\n",
       " 5.24726807192085,\n",
       " 5.76035238493961,\n",
       " 6.3215912463243455,\n",
       " 4.6722955146123155,\n",
       " 5.443278269441182,\n",
       " 5.44173670123618,\n",
       " 4.357409367813885,\n",
       " 3.6746484527343273,\n",
       " 1.0875593004952553,\n",
       " 6.0338733071743,\n",
       " 1.0506717559083876,\n",
       " 1.0694709288043618,\n",
       " 6.318892172324863,\n",
       " 4.669135599893494,\n",
       " 6.0466046360259975,\n",
       " 3.33684447626973,\n",
       " 6.1684671769576145,\n",
       " 1.1632863297567488,\n",
       " 5.9958978999248504,\n",
       " 4.0360774152817145,\n",
       " 1.054496166069463,\n",
       " 1.036476592298914,\n",
       " 5.951851643745354,\n",
       " 6.026289401233958,\n",
       " 1.0643474986152575,\n",
       " 4.033806454974073,\n",
       " 1.0661383667777287,\n",
       " 1.0646924605780648,\n",
       " 3.333506021097399,\n",
       " 1.0491117793786489,\n",
       " 1.6541848874029832,\n",
       " 1.0469768509651558,\n",
       " 5.901546160912221,\n",
       " 4.353454183270078,\n",
       " 1.032293168328919,\n",
       " 5.898772466913232,\n",
       " 1.0607581058580655,\n",
       " 1.2952684592231802,\n",
       " 6.316722814092454,\n",
       " 6.029537288615767,\n",
       " 1.0576187609094758,\n",
       " 1.1593088163026315,\n",
       " 5.43460975753772,\n",
       " 1.154868291629651,\n",
       " 1.043825459957121,\n",
       " 6.164016511467095,\n",
       " 1.4484566114969362,\n",
       " 5.647645131664437,\n",
       " 3.6693368011752545,\n",
       " 4.951365504837914,\n",
       " 3.666164411987071,\n",
       " 1.0754136311361413,\n",
       " 6.0192089347306545,\n",
       " 1.0466033464975721,\n",
       " 5.892313718388283,\n",
       " 5.644753126198445,\n",
       " 4.948852541865603,\n",
       " 1.0584971417357096,\n",
       " 4.348535677294113,\n",
       " 2.3336082011934156,\n",
       " 4.029008124872083,\n",
       " 3.660205040225713,\n",
       " 1.0571433790603126,\n",
       " 1.4440752592760622,\n",
       " 2.3300879432330075,\n",
       " 1.026830920725607,\n",
       " 2.6400950851713474,\n",
       " 1.0534156196386777,\n",
       " 6.122670021493016,\n",
       " 5.637325085975555,\n",
       " 5.241663967897395,\n",
       " 5.945067609066717,\n",
       " 1.0228573648745347,\n",
       " 1.2895513168502526,\n",
       " 4.663731966948084,\n",
       " 5.238608038314439,\n",
       " 5.754309390279219,\n",
       " 1.0851419835862466,\n",
       " 3.6537488500033715,\n",
       " 1.649879497126331,\n",
       " 4.943838739732126,\n",
       " 1.0167485351502235,\n",
       " 2.3250512590048764,\n",
       " 1.0813581229136429,\n",
       " 1.0494175345481176,\n",
       " 5.428484292834579,\n",
       " 6.037930593671599,\n",
       " 6.035740230670267,\n",
       " 1.0457712533449717,\n",
       " 3.6496980006944537,\n",
       " 1.0795937221968324,\n",
       " 1.0438928191400814,\n",
       " 1.0415398951832777,\n",
       " 6.031811327852953,\n",
       " 3.328000619614647,\n",
       " 1.0407326005397295,\n",
       " 6.160588801040324,\n",
       " 1.0411246739723488,\n",
       " 2.31994126920703,\n",
       " 5.749175990294351,\n",
       " 1.1458443983374047,\n",
       " 1.0110452068627296,\n",
       " 6.157751654150539,\n",
       " 1.0348153996862188,\n",
       " 1.2865828091383047,\n",
       " 1.0735861054013958,\n",
       " 1.0388503740453061,\n",
       " 1.282576333371881,\n",
       " 1.0573159601290176,\n",
       " 5.938868452082766,\n",
       " 4.659531603200645,\n",
       " 6.025214706312584,\n",
       " 6.02342417183987,\n",
       " 1.6463214284012524,\n",
       " 2.008467291698564,\n",
       " 4.656432088331169,\n",
       " 5.630906674810788,\n",
       " 4.651869790091143,\n",
       " 1.0548452512053423,\n",
       " 5.626481050442674,\n",
       " 6.018092050065453,\n",
       " 3.642646419305471,\n",
       " 1.0353733387296598,\n",
       " 1.0278806948015826,\n",
       " 5.7432575593600435,\n",
       " 4.339554671179514,\n",
       " 3.638076197558747,\n",
       " 1.0223211638969718,\n",
       " 2.0055604958886213,\n",
       " 1.6397358761655145,\n",
       " 2.9810084757136366,\n",
       " 5.932374878785128,\n",
       " 1.0332766189121145,\n",
       " 4.939948040318402,\n",
       " 1.0305916588816686,\n",
       " 1.0671335008496945,\n",
       " 1.2786040541803056,\n",
       " 1.0512354351890376,\n",
       " 2.0027612016786858,\n",
       " 4.646836699695819,\n",
       " 5.420330100204995,\n",
       " 1.1372565810408706,\n",
       " 1.069882825923493,\n",
       " 2.975648231787336,\n",
       " 3.323146774794343,\n",
       " 6.314468798034563,\n",
       " 5.4162028678256116,\n",
       " 6.1511185339051595,\n",
       " 6.148304615505418,\n",
       " 6.11903192042,\n",
       " 3.3202893359759136,\n",
       " 6.310072611590054,\n",
       " 1.0285841857343627,\n",
       " 1.0223375457946509,\n",
       " 5.231738015028972,\n",
       " 1.274714821207286,\n",
       " 2.3159556197863185,\n",
       " 1.6331305670527674,\n",
       " 1.0136863322456395,\n",
       " 1.9971459141556447,\n",
       " 1.0014113439811445,\n",
       " 2.3137454505033204,\n",
       " 1.0090473688597843,\n",
       " 5.409065655860874,\n",
       " 5.926165161949343,\n",
       " 1.0006630459372439,\n",
       " 3.6289453219356607,\n",
       " 6.00361353514466,\n",
       " 1.9939745483119504,\n",
       " 0.995100761171533,\n",
       " 1.0057678968302681,\n",
       " 1.0438549436942999,\n",
       " 4.332861581528223,\n",
       " 3.624654312298908,\n",
       " 1.4375223921258116,\n",
       " 6.017260247405101,\n",
       " 4.0213245526932795,\n",
       " 5.404183656420396,\n",
       " 1.627050661216377,\n",
       " 6.0137694000207125,\n",
       " 2.306466888448518,\n",
       " 1.0636114868635327,\n",
       " 4.328727183287439,\n",
       " 4.933491359735385,\n",
       " 3.312574350491835,\n",
       " 1.0577278296591244,\n",
       " 2.969349529735486,\n",
       " 1.074014823691796,\n",
       " 1.020502428790641,\n",
       " 1.0544363169668713,\n",
       " 1.0468494518260905,\n",
       " 0.9936312395746143,\n",
       " 1.0521783875471398,\n",
       " 1.2704468677509813,\n",
       " 4.641149468446668,\n",
       " 5.6199204905217774,\n",
       " 1.4303611158330705,\n",
       " 5.987741214934138,\n",
       " 1.4277021288336633,\n",
       " 5.885721900909731,\n",
       " 6.115204118294887,\n",
       " 2.3005103439629946,\n",
       " 5.982873867115239,\n",
       " 1.0145869333030695,\n",
       " 5.61509953137252,\n",
       " 1.0496944873967176,\n",
       " 5.7319131049395375,\n",
       " 0.9910831789203257,\n",
       " 1.0265994921994104,\n",
       " 1.1276058916559026,\n",
       " 5.609248762309003,\n",
       " 4.63507425537882,\n",
       " 1.0452180772908182,\n",
       " 2.962226931920152,\n",
       " 1.0369387658545093,\n",
       " 2.6340818881373163,\n",
       " 5.991804418948142,\n",
       " 2.630251746625825,\n",
       " 4.6293497630157905,\n",
       " 4.926890396848048,\n",
       " 1.011207658618753,\n",
       " 1.0590169213875222,\n",
       " 1.0543308507279558,\n",
       " 4.922987811499583,\n",
       " 6.110699960030501,\n",
       " 0.9860582009016522,\n",
       " 5.8812938964329735,\n",
       " 1.0075901005731218,\n",
       " 1.0029892454551625,\n",
       " 6.107041556478114,\n",
       " 5.9758005028039305,\n",
       " 1.0306440362473963,\n",
       " 1.2640182348534612,\n",
       " 0.984177837394662,\n",
       " 1.4178891412139802,\n",
       " 6.138494614820069,\n",
       " 6.001926239675218,\n",
       " 3.3060368147914416,\n",
       " 6.009460261288169,\n",
       " 6.005285129620355,\n",
       " 1.0059190956474775,\n",
       " 4.31803596023853,\n",
       " 1.0338306780736901,\n",
       " 1.6179504728520446,\n",
       " 1.1196551712663132,\n",
       " 1.6141459547462167,\n",
       " 1.0262065866003245,\n",
       " 5.229222727148426,\n",
       " 5.7220107710035935,\n",
       " 1.2588781517353078,\n",
       " 0.9778679558862291,\n",
       " 1.023866138551614,\n",
       " 1.4133112839698028,\n",
       " 2.2926297961700977,\n",
       " 1.0141829581417459,\n",
       " 4.915451309746466,\n",
       " 6.305794930044868,\n",
       " 5.996815074557735,\n",
       " 5.715126938357609,\n",
       " 3.3004962771497395,\n",
       " 2.621784081252483,\n",
       " 1.0449955551017174,\n",
       " 5.990873000968117,\n",
       " 1.9871366719769186,\n",
       " 1.2542135281548654,\n",
       " 1.9841665033772269,\n",
       " 3.2958320988899112,\n",
       " 1.0389312117544465,\n",
       " 2.616633982526095,\n",
       " 5.9948254139005055,\n",
       " 5.915754708149378,\n",
       " 0.9953555109364497,\n",
       " 1.0366989623748648,\n",
       " 6.09984417028082,\n",
       " 4.309890927689144,\n",
       " 5.967797485183587,\n",
       " 2.610244130566599,\n",
       " 2.2862218246854797,\n",
       " 1.0380466446444485,\n",
       " 1.0263358846397423,\n",
       " 2.605407290370598,\n",
       " 0.9901532755990515,\n",
       " 1.0196897728232268,\n",
       " 2.9504420020992024,\n",
       " 5.989224259694043,\n",
       " 1.2494980586514246,\n",
       " 4.304116993760423,\n",
       " 5.97890933471781,\n",
       " 6.133000490054402,\n",
       " 2.279628182484053,\n",
       " 1.6037753107810118,\n",
       " 0.9700000578251532,\n",
       " 1.0246912330300044,\n",
       " 5.702361058391985,\n",
       " 1.1082829965820031,\n",
       " 5.907165292061586,\n",
       " 4.015454459159489,\n",
       " 4.011133928045795,\n",
       " 5.982072806561926,\n",
       " 0.9622477360276135,\n",
       " 5.977976860218039,\n",
       " 1.0181329361931541,\n",
       " 2.273074935292556,\n",
       " 5.899830795781509,\n",
       " 1.0298275483820247,\n",
       " 5.597891349526412,\n",
       " 1.2435103312030589,\n",
       " 0.9899875603996365,\n",
       " 5.873928500440683,\n",
       " 4.905387403976425,\n",
       " 1.0392009875212285,\n",
       " 1.0336442556796377,\n",
       " 2.266918827199211,\n",
       " 1.05911276082806,\n",
       " 5.224001749009525,\n",
       " 1.039108256633549,\n",
       " 4.624722431716756,\n",
       " 1.5958556710465637,\n",
       " 1.0991163038322624,\n",
       " 4.001663333699357,\n",
       " 0.9825455226974659,\n",
       " 4.2951105757520995,\n",
       " 5.958224381676944,\n",
       " 5.967384459124919,\n",
       " 5.982727047603991,\n",
       " 1.978404154066257,\n",
       " 5.961739108871231,\n",
       " 6.301309184713262,\n",
       " 1.0036994038648186,\n",
       " 0.9765289170123953,\n",
       " 1.0098855620593727,\n",
       " 2.2587247701390245,\n",
       " 5.866498127532843,\n",
       " 5.589526156810392,\n",
       " 1.091314594594769,\n",
       " 3.602801135658681,\n",
       " 3.9931515864776994,\n",
       " 5.951437062674577,\n",
       " 0.9685623505154929,\n",
       " 1.010759335882883,\n",
       " 0.9541137775678208,\n",
       " 4.617065602665036,\n",
       " 1.3991076037295396,\n",
       " 5.859165851246343,\n",
       " 3.9864171223628273,\n",
       " 0.9650960355846393,\n",
       " 6.293571210180632,\n",
       " 0.9709787476262992,\n",
       " 4.611316948837897,\n",
       " 0.9971581226735062,\n",
       " 4.605170326311019,\n",
       " 3.594029876881433,\n",
       " 1.0231085957809656,\n",
       " 2.9392791540075915,\n",
       " 5.3866262160563805,\n",
       " 5.851879279093493,\n",
       " 5.846335182419407,\n",
       " 1.0174711960470806,\n",
       " 1.0024141524811465,\n",
       " 1.0493782509495637,\n",
       " 0.9881806360962101,\n",
       " 0.9446650591077221,\n",
       " 1.0241179328837346,\n",
       " 1.0009204449606803,\n",
       " 6.120259641387998,\n",
       " 6.116343007445007,\n",
       " 0.9388075444488679,\n",
       " 1.0334372020104325,\n",
       " 5.216150900861607,\n",
       " 0.9534383908067549,\n",
       " 5.212959724521395,\n",
       " 6.110395369434123,\n",
       " 1.9724436601340567,\n",
       " 1.3910382594466764,\n",
       " 1.009468969255745,\n",
       " 4.593783878877886,\n",
       " 0.946565734352497,\n",
       " 6.103595356641221,\n",
       " 0.9309651775092234,\n",
       " 3.9743558215536408,\n",
       " 2.594427264328818,\n",
       " 1.0135557255479308,\n",
       " 2.589776300245372,\n",
       " 5.95803838263024,\n",
       " 0.992921581184435,\n",
       " 1.0772331772986907,\n",
       " 5.579015516216339,\n",
       " 1.025025737725588,\n",
       " 5.20388642053925,\n",
       " 0.9863398055831635,\n",
       " 2.2466611723988583,\n",
       " 5.977082392664635,\n",
       " 0.9730439089314792,\n",
       " 0.9772222947922146,\n",
       " 0.9581624817358088,\n",
       " 0.9380640403380974,\n",
       " 5.685916101619663,\n",
       " 5.67850043998287,\n",
       " 5.374867415219435,\n",
       " 5.196592512043451,\n",
       " 3.2796620906261644,\n",
       " 1.2311959647604236,\n",
       " 1.2283081207754971,\n",
       " 1.017263519844322,\n",
       " 5.571909586257521,\n",
       " 5.942818121148059,\n",
       " 4.896165930744341,\n",
       " 0.9785980291497661,\n",
       " 4.282128864378362,\n",
       " 3.2727964185946514,\n",
       " 3.578158172225619,\n",
       " 1.0126441159396855,\n",
       " 2.927990444603423,\n",
       " 0.9730217504432048,\n",
       " 0.9663216872818762,\n",
       " 6.0931692621487805,\n",
       " 1.9661328037436605,\n",
       " 6.090623622442891,\n",
       " 2.237525884612933,\n",
       " 4.889347896055732,\n",
       " 2.9215858311099785,\n",
       " 6.085469955000118,\n",
       " 5.934358038866775,\n",
       " 4.275293950631472,\n",
       " 5.944182395633106,\n",
       " 6.079643987009728,\n",
       " 5.877108210932856,\n",
       " 5.66208116539973,\n",
       " 3.5680100368143814,\n",
       " 5.936709716770853,\n",
       " 5.935977267105267,\n",
       " 6.282545911091098,\n",
       " 6.072552984193935,\n",
       " 3.2640321399752,\n",
       " 5.364759977341526,\n",
       " 5.359799289300605,\n",
       " 0.9486164808639364,\n",
       " 2.227635296033236,\n",
       " 4.581215682333866,\n",
       " 4.87889494437793,\n",
       " 2.9127669356076256,\n",
       " 1.0132475758360586,\n",
       " 5.967064319876651,\n",
       " 6.0842911787302745,\n",
       " 0.987775344706355,\n",
       " 1.0076164120719033,\n",
       " 4.872442143843364,\n",
       " 5.561783891726933,\n",
       " 1.2166118624997153,\n",
       " 1.0021943687760049,\n",
       " 1.2115870910581585,\n",
       " 5.9222683951361486,\n",
       " 6.063822575460759,\n",
       " 3.255473931849624,\n",
       " 5.647967535107597,\n",
       " 4.862742842361304,\n",
       " 0.9953902113125592,\n",
       " 5.63978639603236,\n",
       " 0.9995768950390227,\n",
       " 2.5747249740468523,\n",
       " 0.9122167804004158,\n",
       " 6.076452278101217,\n",
       " 3.2478687608672607,\n",
       " 2.567516262893045,\n",
       " 5.956826235975168,\n",
       " 0.9060381452796343,\n",
       " 0.9232847637947268,\n",
       " 2.901818574519807,\n",
       " 2.8956664467629323,\n",
       " 1.0357037808155005,\n",
       " 1.2036713287044705,\n",
       " 5.910158357389837,\n",
       " 0.9175073785368937,\n",
       " 5.924008877843095,\n",
       " 0.9859523089794111,\n",
       " 5.857994508572005,\n",
       " 5.949298587561765,\n",
       " 1.3733614432842942,\n",
       " 5.849660781023495,\n",
       " 0.9519862806794571,\n",
       " 5.622642130504255,\n",
       " 5.34429430309127,\n",
       " 5.841658572326439,\n",
       " 6.053777651823882,\n",
       " 3.238383364965761,\n",
       " 1.0036607543087637,\n",
       " 6.066207887584497,\n",
       " 5.612691990879059,\n",
       " 5.604848158953991,\n",
       " 0.9918472738129503,\n",
       " 2.5560090733991427,\n",
       " 1.364188220253926,\n",
       " 5.8263930730304745,\n",
       " 0.9840077890344322,\n",
       " 0.9962727003502525,\n",
       " 6.046313252937107,\n",
       " 5.589670673702682,\n",
       " 1.1932643941685668,\n",
       " 0.8907672851157903,\n",
       " 5.334514679173047,\n",
       " 0.9065464493706719,\n",
       " 5.912966987228784,\n",
       " 0.974403644276184,\n",
       " 1.0225329153949547,\n",
       " 5.19155204906921,\n",
       " 6.038703416461152,\n",
       " 1.187576407935335,\n",
       " 3.9535767442440033,\n",
       " 6.032070922682569,\n",
       " 0.9674021988823821,\n",
       " 0.9864155411284004,\n",
       " 0.9347038977324547,\n",
       " 0.8823548750695359,\n",
       " 0.9591704906146336,\n",
       " 2.5438776373446492,\n",
       " 0.8970092087213555,\n",
       " 1.1803277817050366,\n",
       " 5.570712038855276,\n",
       " 5.8340032158622614,\n",
       " 5.892076800113112,\n",
       " 0.976613851405366,\n",
       " 1.173944086940199,\n",
       " 2.875855752379119,\n",
       " 5.1833115595956665,\n",
       " 1.9576798684473848,\n",
       " 0.9380469051463355,\n",
       " 0.973828342328533,\n",
       " 0.8709607608081217,\n",
       " 0.9789525164136842,\n",
       " 1.5675015038961637,\n",
       " 0.8636908550410006,\n",
       " 0.9519747548771736,\n",
       " 5.319782092597489,\n",
       " 1.1649485388593053,\n",
       " 1.01099203733385,\n",
       " 0.9657756393941789,\n",
       " 5.879918833019295,\n",
       " 1.5579694452504218,\n",
       " 0.9433231268547879,\n",
       " 5.546174953882359,\n",
       " 0.988093295285195,\n",
       " 5.1719565147710425,\n",
       " 2.8650476807226046,\n",
       " 5.80034786602311,\n",
       " 3.536067979213583,\n",
       " 1.1577261872188043,\n",
       " 1.9495196272205428,\n",
       " 5.538312775818222,\n",
       " 1.547706193650774,\n",
       " 0.9694974939583698,\n",
       " 5.789356773938578,\n",
       " 6.0191693999865095,\n",
       " 2.5321029012391527,\n",
       " 1.1499023354757156,\n",
       " 6.01178305229344,\n",
       " 5.528971894190286,\n",
       " 5.164922323327298,\n",
       " 5.896147561225644,\n",
       " 5.157280997368847,\n",
       " 0.9326026216279805,\n",
       " 1.142640030939047,\n",
       " 3.524467519716878,\n",
       " 4.25003516063827,\n",
       " 6.003554201369129,\n",
       " 0.9984334214100843,\n",
       " 5.822926460539234,\n",
       " 3.937894899965455,\n",
       " 4.842409363228271,\n",
       " 5.996305366812795,\n",
       " 4.834618787004243,\n",
       " 0.9473036886673584,\n",
       " 5.14732141339119,\n",
       " 4.8252024796671,\n",
       " 1.0361557676344892,\n",
       " 2.5202560245475634,\n",
       " 5.896380233348913,\n",
       " 4.814947650438155,\n",
       " 2.851635327742518,\n",
       " 6.044453614136334,\n",
       " 5.543986339237877,\n",
       " 4.5666069243945415,\n",
       " 1.1304640310741685,\n",
       " 5.302771311443089,\n",
       " 3.5088956223376653,\n",
       " 0.9210905482526697,\n",
       " 4.800708726023045,\n",
       " 1.0267014452029153,\n",
       " 0.9206415118154869,\n",
       " 0.8770361707107674,\n",
       " 5.516505061780215,\n",
       " 4.556061007012711,\n",
       " 5.880605878970804,\n",
       " 5.810158812374619,\n",
       " 0.9129054590102001,\n",
       " 0.9342764280143717,\n",
       " 6.27544819824655,\n",
       " 2.5083821466449927,\n",
       " 0.9627310820482743,\n",
       " ...]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def train(diffusion_model, optimizer, batch_size, nb_epochs, device):\n",
    "\n",
    "\n",
    "    training_loss = []\n",
    "\n",
    "    for epoch in tqdm(range(nb_epochs)):\n",
    "\n",
    "        x0 = sample_batch(batch_size).to(device)\n",
    "        loss = diffusion_model.get_loss(x0)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        training_loss.append(loss.item())\n",
    "\n",
    "        if epoch % 5000 == 0:\n",
    "            plt.plot(training_loss)\n",
    "            plt.savefig(f'figs/training_loss_epoch_{epoch}.png')\n",
    "            plt.close()\n",
    "\n",
    "            plot(diffusion_model, f'figs/training_epoch_{epoch}.png', device)\n",
    "\n",
    "    return training_loss\n",
    "\n",
    "\n",
    "device = 'cuda'\n",
    "mlp_model = MLP(hidden_dim=128).to(device)\n",
    "model = DiffusionModel(40, mlp_model, device)\n",
    "optimizer = torch.optim.Adam(mlp_model.parameters(), lr=1e-4)\n",
    "\n",
    "train(model, optimizer, 64_000, 300_000, device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}