{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data2/packages/anaconda3/envs/ebm/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.distributions import MultivariateNormal\n",
    "torch.cuda.set_device(1)\n",
    "\n",
    "class ToyDataset(object):\n",
    "    def __init__(self, data_name):\n",
    "        self.data_name = data_name\n",
    "        self.rng = np.random.RandomState()\n",
    "\n",
    "    def gen_batch(self, batch_size):\n",
    "        return inf_train_gen(self.data_name, self.rng, batch_size)\n",
    "    \n",
    "    def log_prob(self, x):\n",
    "        if self.data_name == '25gaussians':\n",
    "            return log_prob_25gaussians(x)\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "        \n",
    "def inf_train_gen(data, rng=None, batch_size=200):\n",
    "    if data == \"25gaussians\":\n",
    "        scale = 4.\n",
    "        centers = [(x/2, y/2) for x in range(-2,3) for y in range(-2,3)]\n",
    "        centers = [(scale * x, scale * y) for x, y in centers]\n",
    "\n",
    "        dataset = []\n",
    "        for i in range(batch_size):\n",
    "            point = rng.randn(2) * 0.2\n",
    "            idx = rng.randint(25)\n",
    "            center = centers[idx]\n",
    "            point[0] += center[0]\n",
    "            point[1] += center[1]\n",
    "            dataset.append(point)\n",
    "        dataset = np.array(dataset, dtype=\"float32\")\n",
    "        dataset /= 1.414\n",
    "        return dataset\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "\n",
    "def log_prob_25gaussians(x):\n",
    "    scale = 4. / 1.414\n",
    "    centers = [(x/2, y/2) for x in range(-2,3) for y in range(-2,3)]\n",
    "    centers = [(scale * x, scale * y) for x, y in centers]\n",
    "\n",
    "    conv = (1 / 25 / (1.414**2) * torch.eye(2)).to(x.device)\n",
    "    gaussians = [MultivariateNormal(loc=torch.tensor(mu, device=x.device), covariance_matrix=conv) for mu in centers]\n",
    "\n",
    "    log_probs = torch.stack([g.log_prob(x).unsqueeze(1) for g in gaussians], dim=1)\n",
    "    log_probs = torch.logsumexp(log_probs, dim=1) - np.log(25.) #- np.log(2*np.pi) - 2*np.log(5) - 2*np.log(1.414)\n",
    "    return log_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Models and Training Algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.autograd as autograd\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "class ToyModel(nn.Module):\n",
    "    def __init__(self, n_output=1):\n",
    "        super().__init__()\n",
    "        self.lin1 = nn.Linear(2, 128)\n",
    "        self.lin2 = nn.Linear(128, 128)\n",
    "        self.lin3 = nn.Linear(128, 128)\n",
    "        self.lin4 = nn.Linear(128, 128)\n",
    "        self.lin5 = nn.Linear(128, n_output)\n",
    "\n",
    "        self.act_func = nn.Softplus()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.act_func(self.lin1(x))\n",
    "        x = self.act_func(self.lin2(x))\n",
    "        x = self.act_func(self.lin3(x))\n",
    "        x = self.act_func(self.lin4(x))\n",
    "        x = self.lin5(x)\n",
    "        return x\n",
    "    \n",
    "def energy_discrepancy(energy_net, samples, m_particles=4, t_noise=1.0, w_stable=1.0):\n",
    "    device = samples.device\n",
    "\n",
    "    browniani = torch.randn_like(samples).to(device) * t_noise\n",
    "    brownianij = torch.randn(samples.size(0), m_particles, *samples.shape[1:]).to(device) * t_noise\n",
    "\n",
    "    pert_data = samples.unsqueeze(1) + browniani.unsqueeze(1) + brownianij\n",
    "\n",
    "    pos_energy = energy_net(samples)\n",
    "    neg_energy = energy_net(pert_data.view(-1, *samples.shape[1:])).view(samples.size(0), -1)\n",
    "    val = pos_energy - neg_energy\n",
    "    if w_stable != 0:\n",
    "        val = torch.cat([val, np.log(w_stable) * torch.ones_like(val[:, :1])], dim=-1)\n",
    "\n",
    "    loss = val.logsumexp(dim=-1).mean()\n",
    "    return loss\n",
    "\n",
    "def score_matching(energy_net, samples):\n",
    "    samples.requires_grad_(True)\n",
    "    logp = -energy_net(samples).sum()\n",
    "    grad1 = autograd.grad(logp, samples, create_graph=True)[0]\n",
    "    loss1 = torch.norm(grad1, dim=-1) ** 2 / 2.\n",
    "    loss2 = torch.zeros(samples.shape[0], device=samples.device)\n",
    "\n",
    "    for i in range(samples.shape[1]):\n",
    "        grad = autograd.grad(grad1[:, i].sum(), samples, create_graph=True, retain_graph=True)[0][:, i]\n",
    "        loss2 += grad\n",
    "\n",
    "    loss = loss1 + loss2\n",
    "\n",
    "    return loss.mean()\n",
    "\n",
    "def contrastive_divergence(energy_net, samples):\n",
    "    num_steps = 10\n",
    "    step_size = 0.01\n",
    "    noise = torch.autograd.Variable(samples.clone(), requires_grad=True)\n",
    "    for _ in range(num_steps):\n",
    "        log_prob = -energy_net(noise)\n",
    "        grad_y = torch.autograd.grad(log_prob.sum(), [noise], retain_graph=False)[0]\n",
    "        noise.data += 0.5 * step_size * grad_y + np.sqrt(step_size) * torch.randn_like(noise)\n",
    "    samples_neg = noise.detach()\n",
    "\n",
    "    log_p_x_pos = -energy_net(samples)\n",
    "    log_p_x_neg = -energy_net(samples_neg)\n",
    "\n",
    "    loss_cd = torch.mean(log_p_x_neg - log_p_x_pos)\n",
    "    loss_reg = torch.mean(log_p_x_pos**2 + log_p_x_neg**2)\n",
    "    loss = loss_cd + .3 * loss_reg\n",
    "    return loss\n",
    "\n",
    "def mse_log_prob(dataset, energy_net):\n",
    "    samples = dataset.gen_batch(5000)\n",
    "    samples = torch.from_numpy(samples).float()\n",
    "\n",
    "    if device is not None:\n",
    "        samples = samples.to(device)\n",
    "\n",
    "    log_ground_truth = dataset.log_prob(samples)\n",
    "\n",
    "    log_prob = -energy_net(samples)\n",
    "    logZ = torch.logsumexp(-energy_net(samples)-log_ground_truth, dim=(0)) - np.log(5000)\n",
    "    log_prob_ebm = log_prob - logZ\n",
    "\n",
    "    return torch.mean((log_prob_ebm - log_ground_truth) ** 2).cpu().item(), logZ.cpu().item()\n",
    "\n",
    "def train(model, ema_model, optimizer, dataset, compute_mse=True, algo='ed', m_particles=4, t_noise=1.0, w_stable=1.0):\n",
    "    mse_list = []\n",
    "    for step in tqdm(range(50001)):\n",
    "        samples = dataset.gen_batch(128)\n",
    "        samples = torch.from_numpy(samples).float().to(device)\n",
    "        \n",
    "        if algo == 'ed':\n",
    "            loss = energy_discrepancy(model, samples, m_particles, t_noise, w_stable)\n",
    "        elif algo == 'sm':\n",
    "            loss = score_matching(model, samples)\n",
    "        elif algo == 'cd':\n",
    "            loss = contrastive_divergence(model, samples)\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        with torch.no_grad():\n",
    "            for p, ema_p in zip(model.parameters(), ema_model.parameters()):\n",
    "                ema_p.mul_(0.999).add_(p, alpha=0.001)\n",
    "\n",
    "        if compute_mse and step % 100 == 0:\n",
    "            mse_list.append(mse_log_prob(dataset, ema_model)[0])\n",
    "\n",
    "    if compute_mse:\n",
    "        return np.array(mse_list)\n",
    "    else:\n",
    "        return\n",
    "\n",
    "def visualize_energynet(energy_net, title, left_bound=-4., right_bound=4.):\n",
    "    mesh = []\n",
    "    grid_size = 100\n",
    "    x = torch.linspace(left_bound, right_bound, grid_size)\n",
    "    y = torch.linspace(left_bound, right_bound, grid_size)\n",
    "    xx, yy = torch.meshgrid(x, y, indexing='ij')\n",
    "    mesh = torch.stack([xx, yy], dim=-1).permute(1, 0, 2).reshape(-1, 2)\n",
    "    mesh = mesh.to(device)\n",
    "\n",
    "    prob = -energy_net(mesh).view(grid_size, grid_size)\n",
    "    logZ = torch.logsumexp(prob, dim=(0, 1))\n",
    "    prob = torch.exp(prob - logZ)\n",
    "\n",
    "    plt.figure(figsize=(2,2))\n",
    "    plt.imshow(np.flipud(prob.detach().cpu().numpy()), cmap='viridis', extent=[left_bound, right_bound, left_bound, right_bound])\n",
    "    plt.title(title)\n",
    "    plt.axis('off')\n",
    "    plt.axis('square')\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "def plot_mse(mse_list, title):\n",
    "    x = np.array([0, 100, 200, 300, 400, 500])\n",
    "    plt.figure(figsize=(5,3.3), tight_layout=True)\n",
    "    plt.plot(mse_list, linewidth=1.)\n",
    "    plt.xticks(x, ['0', '10k', '20k', '30k', '40k', '50k'])\n",
    "    plt.title(title, fontsize=12)\n",
    "    plt.xlabel('Iterations', fontsize=12)\n",
    "    plt.ylabel('MSE of log-density', fontsize=12)\n",
    "    plt.ylim([-0.2, 5.2])\n",
    "    plt.grid(linestyle='--', linewidth=0.5, alpha=0.5)\n",
    "    ax=plt.gca()\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    plt.show()\n",
    "    plt.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparing ED with CD & SM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:59<00:00, 167.09it/s]\n",
      "100%|██████████| 50001/50001 [12:23<00:00, 67.25it/s]\n",
      "100%|██████████| 50001/50001 [13:17<00:00, 62.70it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23076 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 37704 characters]",
      "text/plain": [
       "<Figure size 500x330 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 11292 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48048 characters]",
      "text/plain": [
       "<Figure size 500x330 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 40480 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53552 characters]",
      "text/plain": [
       "<Figure size 500x330 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "dataset = ToyDataset('25gaussians')\n",
    "\n",
    "ebm_ed = ToyModel().to(device)\n",
    "optimizer_ed = optim.Adam(ebm_ed.parameters(), lr=1e-3)\n",
    "ema_ebm_ed = ToyModel(n_output=1).to(device)\n",
    "ema_ebm_ed.load_state_dict(ebm_ed.state_dict())\n",
    "\n",
    "ebm_sm = ToyModel().to(device)\n",
    "optimizer_sm = optim.Adam(ebm_sm.parameters(), lr=1e-3)\n",
    "ema_ebm_sm = ToyModel(n_output=1).to(device)\n",
    "ema_ebm_sm.load_state_dict(ebm_sm.state_dict())\n",
    "\n",
    "ebm_cd = ToyModel().to(device)\n",
    "optimizer_cd = optim.Adam(ebm_cd.parameters(), lr=1e-3)\n",
    "ema_ebm_cd = ToyModel(n_output=1).to(device)\n",
    "ema_ebm_cd.load_state_dict(ebm_cd.state_dict())\n",
    "\n",
    "mes_ed = train(ebm_ed, ema_ebm_ed, optimizer_ed, dataset, algo='ed', m_particles=4, t_noise=1.0, w_stable=1.0)\n",
    "mes_sm = train(ebm_sm, ema_ebm_sm, optimizer_sm, dataset, algo='sm')\n",
    "mes_cd = train(ebm_cd, ema_ebm_cd, optimizer_cd, dataset, algo='cd')\n",
    "\n",
    "visualize_energynet(ema_ebm_ed, 'Energy Discrepancy')\n",
    "plot_mse(mes_ed, 'Energy Discrepancy')\n",
    "visualize_energynet(ema_ebm_sm, 'Score Matching')\n",
    "plot_mse(mes_sm, 'Score Matching')\n",
    "visualize_energynet(ema_ebm_cd, 'Contrastive Divergence')\n",
    "plot_mse(mes_cd, 'Contrastive Divergence')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qualitatively Understanding the Effect of $t$ and $M$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$t=0.05, M=1, 32, 512, w=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:34<00:00, 182.41it/s]\n",
      "100%|██████████| 50001/50001 [04:42<00:00, 176.86it/s]\n",
      "100%|██████████| 50001/50001 [05:48<00:00, 143.67it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 5420 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6336 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 12012 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm1 = ToyModel().to(device)\n",
    "optimizer_ebm1 = optim.Adam(ebm1.parameters(), lr=1e-3)\n",
    "ema_ebm1 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm1.load_state_dict(ebm1.state_dict())\n",
    "\n",
    "ebm2 = ToyModel().to(device)\n",
    "optimizer_ebm2 = optim.Adam(ebm2.parameters(), lr=1e-3)\n",
    "ema_ebm2 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm2.load_state_dict(ebm2.state_dict())\n",
    "\n",
    "ebm3 = ToyModel().to(device)\n",
    "optimizer_ebm3 = optim.Adam(ebm3.parameters(), lr=1e-3)\n",
    "ema_ebm3 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm3.load_state_dict(ebm3.state_dict())\n",
    "\n",
    "train(ebm1, ema_ebm1, optimizer_ebm1, dataset, compute_mse=False, algo='ed', m_particles=1, t_noise=0.05, w_stable=1.0)\n",
    "train(ebm2, ema_ebm2, optimizer_ebm2, dataset, compute_mse=False, algo='ed', m_particles=32, t_noise=0.05, w_stable=1.0)\n",
    "train(ebm3, ema_ebm3, optimizer_ebm3, dataset, compute_mse=False, algo='ed', m_particles=512, t_noise=0.05, w_stable=1.0)\n",
    "\n",
    "visualize_energynet(ema_ebm1, 't=0.05, m=1')\n",
    "visualize_energynet(ema_ebm2, 't=0.05, m=32')\n",
    "visualize_energynet(ema_ebm3, 't=0.05, m=512')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$t=50, M=1, 32, 512, w=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:42<00:00, 177.01it/s]\n",
      "100%|██████████| 50001/50001 [04:46<00:00, 174.33it/s]\n",
      "100%|██████████| 50001/50001 [05:44<00:00, 144.97it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17232 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23860 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22184 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm1 = ToyModel().to(device)\n",
    "optimizer_ebm1 = optim.Adam(ebm1.parameters(), lr=1e-3)\n",
    "ema_ebm1 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm1.load_state_dict(ebm1.state_dict())\n",
    "\n",
    "ebm2 = ToyModel().to(device)\n",
    "optimizer_ebm2 = optim.Adam(ebm2.parameters(), lr=1e-3)\n",
    "ema_ebm2 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm2.load_state_dict(ebm2.state_dict())\n",
    "\n",
    "ebm3 = ToyModel().to(device)\n",
    "optimizer_ebm3 = optim.Adam(ebm3.parameters(), lr=1e-3)\n",
    "ema_ebm3 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm3.load_state_dict(ebm3.state_dict())\n",
    "\n",
    "train(ebm1, ema_ebm1, optimizer_ebm1, dataset, compute_mse=False, algo='ed', m_particles=1, t_noise=50, w_stable=1.0)\n",
    "train(ebm2, ema_ebm2, optimizer_ebm2, dataset, compute_mse=False, algo='ed', m_particles=32, t_noise=50, w_stable=1.0)\n",
    "train(ebm3, ema_ebm3, optimizer_ebm3, dataset, compute_mse=False, algo='ed', m_particles=512, t_noise=50, w_stable=1.0)\n",
    "\n",
    "visualize_energynet(ema_ebm1, 't=50, m=1')\n",
    "visualize_energynet(ema_ebm2, 't=50, m=32')\n",
    "visualize_energynet(ema_ebm3, 't=50, m=512')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qualitatively Understanding the Effect of $w$ and $M$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$t=1, M=1, 32, 512, w=0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:25<00:00, 188.18it/s]\n",
      "100%|██████████| 50001/50001 [04:38<00:00, 179.73it/s]\n",
      "100%|██████████| 50001/50001 [05:42<00:00, 145.90it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAK4AAADECAYAAAAGYxrSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAH8klEQVR4nO3bW0iVXR7H8d9Ws23lIalGIw8XLw2KQpQYiYeophQtSFAkrAgag4yaCLqIurQIy4uKyLqIyC4qCAymcLww3KUwdNAiBqEyqUEKIjyUGtKaC+khp3dedVvt/s73A4Iun8N68Mty6Wb7nHNOgDFhoZ4AEAzChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhDuBtrY25ebmas6cOUpISNDevXs1ODgY6mlNWVdXl/bv36+cnBz5/X75fD69fPky1NMKGuH+gY6ODq1du1YfP35UXV2ddu7cqfPnz6usrCzUU5uy9vZ2nTp1SgMDA0pLSwv1dKbP4X8qKipyiYmJrq+vzxu7cOGCk+SamppCOLOpe/funevv73fOOVdbW+skue7u7tBOahrMrLiPHz+Wz+fTzZs3vbEHDx7I5/Np+fLl444tKirSypUrp3W//v5+NTc3q7KyUjExMd74tm3bNG/ePF27di2o66ampqqkpER37txRVlaWoqKilJmZqTt37kiSbty4oczMTPn9fq1YsUKPHj2a1nN8ER8fr+jo6O9yrV9BRKgnMFkZGRmKi4tTa2urNm3aJEkKBAIKCwtTZ2en+vv7FRMTo8+fP6utrU1VVVXeuYODgxoeHp7wHrNmzVJsbKwk6cmTJxodHVVWVta4YyIjI7Vs2bJpBfXs2TNt2bJFu3btUmVlpU6cOKGNGzfq3LlzOnTokHbv3i1JOnbsmMrLy9XV1aWwsLE1ZmRkRAMDA5O6z4IFC4Ke4y8v1Ev+VBQXF7vs7Gzv69LSUldaWurCw8Pd7du3nXPOPXz40ElyjY2N3nHbt293kib8KCgo8M65fv26k+RaW1u/mUdZWZlLSEgI6hlSUlKcJNfW1uaNNTU1OUkuKirK9fT0eOP19fVOkmtpafHGLl68OKln+aMf7UzYKphZcSUpLy9Phw8f1ocPHzR37lzdvXtXR48eVU9PjwKBgAoLCxUIBOTz+ZSbm+udd/DgQVVWVk54/fnz53ufDw0NSZJmz579zXF+v9/7fjDS09O1atUq7+sv25o1a9YoOTn5m/EXL15o9erVkqQNGzaoubk56HvPFObCHR0dVXt7u5KSkvT27Vvl5eXp6dOnCgQCksa2D+np6YqPj/fOS09PV3p6+pTuFRUVJWnsV/N/Gx4e9r4fjK/jlORtT5KSkn53/P37995YYmKiEhMTg773TGEq3KysLPn9frW2tio5OVmLFi3S0qVLlZeXp7Nnz2pkZESBQECbN28ed15fX9+kVsjIyEgv+C9x9Pb2fnNcb2+vFi9eHPRzhIeHT2ncffXuqqGhIfX19U3qPgkJCVOfnBGmwo2MjFR2drYCgYCSk5OVl5cnaWwlHhkZ0ZUrV/TmzRvl5+ePO2/fvn26dOnShNcvKCjw/rrPyMhQRESE7t+/r/Lycu+YT58+qaOjY9zYz3T16lXt2LFjUse6Gfx2QlPhSmOR1tXV6fnz5zpw4ICksb+e09LSdPz4ce+YrwWzx42NjdW6devU0NCgI0eOeP9Kunz5sgYHB0P2IgR73DEmw62pqdGrV6/GBZqfn6/6+nqlpqZqyZIl484JZo8rSTU1NcrJyVFBQYGqqqr0+vVrnTx5UuvXr1dhYeG4Y30+37gV+0cJdo/b19en06dPS5Lu3bsnSTpz5ozi4uIUFxenPXv2fNd5/nCh/rfGVPX397vw8HAXHR3tRkdHvfGGhgYnyW3duvW73i8QCLicnBzn9/vdwoULXXV1tfcK1BcDAwNOkquoqJjweikpKa64uPibcUmuurp63Fh3d7eT5Gpra6f3EF9d6/c+UlJSpn39n83n3AzeCP0kt27dUklJiTo7O5WZmRnq6fxfMPOS76+spaVFFRUVRPsTseLCJFZcmES4MIlwYRLhwiTChUmTfuXsL2H23mcFe5o/X5/Ucay4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4P0BEwp8U/uffpLDwUE9lxiLcH+BfR1L1t783KmLRglBPZcaKCPUEZqIl/3Da/++/KnmgM9RTmbEI9weIavynljRKn0M9kRmMrQJMIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwiTChUmEC5MIFyYRLkwiXJhEuDCJcGES4cIkwoVJhAuTCBcmES5MIlyYRLgwiXBhEuHCJJ9zzoV6EsBUseLCJMKFSYQLkwgXJhEuTCJcmES4MIlwYRLhwqT/ACg6y02wJYcEAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAK4AAADECAYAAAAGYxrSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAJsElEQVR4nO3cfUiT/R7H8c906TTNVjpUTKs/CoejqGElTkM7ZbfWXZEyyoo4UFBkQVCcMqo/7AFN4kRRRI0e/MOiov7oAYnChUL0/ASdatVJcHmQdC1spf7OH3JftNtKndX2ve/PC4T87Xdd1++it5fXNpxOKaVAJExYsBdAFAiGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIx3C80NDQgOzsb0dHRSExMRFlZGbxeb7CXNWDnzp3DrFmzkJycjMjISKSkpGDhwoV49OiR37zW1lZUVlYiJycHCQkJGD58OKZOnYra2togrbz/9MFeQKi4d+8e8vPzkZ6ejurqajQ1NaGqqgrPnj3DpUuXgr28AXn48CGMRiPWrl2L+Ph4uN1uHD16FJmZmWhsbMSECRMAAI2Njdi8eTN+++03lJeXQ6/X48yZM7Db7Xjy5Am2b98e5DP5DkVKKaVmz56tkpKSVHt7uzZ2+PBhBUBduXIliCv7Mdxut9Lr9WrlypXamMvlUq9evfKb193drfLy8lRkZKTyer2/epn9FpK3Cg8ePIBOp8OFCxe0sdu3b0On02HSpEl+c2fPno0pU6YM6ngejwd1dXUoLS3FsGHDtPGlS5ciJiYGp06dCmi/o0ePRlFREa5fvw6r1YqoqChYLBZcv34dAHD27FlYLBYYDAZMnjwZd+/eHdR5fI/JZEJ0dDTa2tq0sTFjxiAtLc1vnk6nw7x58+Dz+eByuX7aegYrJG8VMjIyMHz4cNTX12Pu3LkAAKfTibCwMNy/fx8ejwfDhg1Dd3c3GhoasGLFCm1br9eLjx8/9nmMIUOGIC4uDkDPr9bOzk5YrVa/OREREZg4ceKggnr+/DkWLVqElStXorS0FFVVVZgzZw4OHjyITZs2YdWqVQCAnTt3oqSkBE+fPkVYWM/1xOfz4f379/06Tnx8fK+xtrY2fP78GW63G3v37oXH40F+fn6f+3K73d/cZ8gI9iX/WwoLC1VmZqb2/YIFC9SCBQtUeHi4unTpklJKqTt37igA6vz589q8ZcuWKQB9fuXm5mrbnD59WgFQ9fX1vdZRXFysEhMTAzqHtLQ0BUA1NDRoY1euXFEAVFRUlHr9+rU2fujQIQVAXbt2TRtzOBz9Opdv/TeOHz9eezwmJkaVl5errq6u7665tbVVmUwmZbPZAjrnXyUkr7gAYLPZUF5ejg8fPmDo0KG4ceMGduzYgdevX8PpdKKgoABOpxM6nQ7Z2dnadhs2bEBpaWmf+zcajdq/Ozo6AACRkZG95hkMBu3xQJjNZkybNk37/o/bmry8PKSmpvYad7lcmD59OgBg1qxZqKurC/jYDocDHo8HLpcLDocDHR0d6Orq0q7of9bd3Y3Fixejra0N+/btC/i4v0JIh9vZ2YnGxkaMGjUKLS0tsNlsePz4MZxOJ4Ce2wez2YwRI0Zo25nNZpjN5gEdKyoqCkDPr+Y/+/jxo/Z4IL6ME4B2ezJq1Kivjr97904bS0pKQlJSUsDH/vIHxm63Iz09HQBQVVX11flr1qzB5cuXcfz4ce2Vh1AVsuFarVYYDAbU19cjNTUVJpMJ48aNg81mw4EDB+Dz+eB0OjF//ny/7drb2/t1hYyIiNCC/yOO5ubmXvOam5uRnJwc8HmEh4cPaFx98ZdUHR0daG9v79dxEhMTv/u40WhEXl4eampqvhru9u3bceDAAezatQtLlizp1zGDKWTDjYiIQGZmJpxOJ1JTU2Gz2QD0XIl9Ph9qamrw9u1b5OTk+G23du1aHDt2rM/95+bmas/uMzIyoNfrcevWLZSUlGhzPn36hHv37vmN/Uq1tbVYvnx5v+aqfvzp4Ld+EPbv349t27Zh3bp12Lhx44DXGQwhGy7QE2l1dTVevHiB9evXA+h5ppueno7du3drc74UyD1uXFwcZsyYgZMnT2LLli2IjY0FAJw4cQJerxfFxcU/6pQGJNB73JaWFphMJr+xV69e4erVq71eOamtrUVZWRkWL16M6urqQa33Vwr5cCsqKvDmzRu/QHNycnDo0CGMHj0aKSkpftsEco8LABUVFcjKykJubi5WrFiBpqYm7NmzBzNnzkRBQYHfXJ1O53fF/lkCvce1WCzIz8/HxIkTYTQa8ezZMxw5cgSfP3/Grl27tHk3b97E0qVLMXLkSOTn56OmpsZvP1lZWRg7duygz+OnCPbLGt/j8XhUeHi4io2NVZ2dndr4yZMnFQC1ZMmSH3o8p9OpsrKylMFgUAkJCWr16tXK4/H4zXn//r0CoOx2e5/7S0tLU4WFhb3GAajVq1f7jb18+VIBUJWVlYM7CaXU1q1bldVqVUajUen1epWcnKzsdrt68OCB37y+Xm5zOByDXsvPolOKn6swEBcvXkRRURHu378Pi8US7OX8bYXkW76h7Nq1a7Db7Yw2yHjFJZF4xSWRGC6JxHBJJIZLIjFcEqnf75z9Iyw4b3vS30td9+l+zeMVl0RiuCQSwyWRGC6JxHBJJIZLIjFcEonhkkgMl0RiuCQSwyWRGC6JxHBJJIZLIjFcEonhkkgMl0RiuCQSwyWRGC6JFJRww6Kj0fSvLHTMywzG4ekvICjh6mJjsOefR/CmoO+5RF8TlA927vpfK/4993eY3/0XncFYAIkXnE8k7+5C15P/BOXQ9NfAJ2ckEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkEsMlkRguicRwSSSGSyIxXBKJ4ZJIDJdEYrgkkk4ppYK9CKKB4hWXRGK4JBLDJZEYLonEcEkkhksiMVwSieGSSAyXRPo/sSAxr3WobGQAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16272 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm1 = ToyModel().to(device)\n",
    "optimizer_ebm1 = optim.Adam(ebm1.parameters(), lr=1e-3)\n",
    "ema_ebm1 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm1.load_state_dict(ebm1.state_dict())\n",
    "\n",
    "ebm2 = ToyModel().to(device)\n",
    "optimizer_ebm2 = optim.Adam(ebm2.parameters(), lr=1e-3)\n",
    "ema_ebm2 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm2.load_state_dict(ebm2.state_dict())\n",
    "\n",
    "ebm3 = ToyModel().to(device)\n",
    "optimizer_ebm3 = optim.Adam(ebm3.parameters(), lr=1e-3)\n",
    "ema_ebm3 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm3.load_state_dict(ebm3.state_dict())\n",
    "\n",
    "train(ebm1, ema_ebm1, optimizer_ebm1, dataset, compute_mse=False, algo='ed', m_particles=1, t_noise=1, w_stable=0)\n",
    "train(ebm2, ema_ebm2, optimizer_ebm2, dataset, compute_mse=False, algo='ed', m_particles=32, t_noise=1, w_stable=0)\n",
    "train(ebm3, ema_ebm3, optimizer_ebm3, dataset, compute_mse=False, algo='ed', m_particles=512, t_noise=1, w_stable=0)\n",
    "\n",
    "visualize_energynet(ema_ebm1, 'w=0, m=1')\n",
    "visualize_energynet(ema_ebm2, 'w=0, m=32')\n",
    "visualize_energynet(ema_ebm3, 'w=0, m=512')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$t=1, M=1, 32, 512, w=0.05$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:34<00:00, 182.32it/s]\n",
      "100%|██████████| 50001/50001 [04:42<00:00, 176.80it/s]\n",
      "100%|██████████| 50001/50001 [05:44<00:00, 145.13it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 11228 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16924 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21880 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm1 = ToyModel().to(device)\n",
    "optimizer_ebm1 = optim.Adam(ebm1.parameters(), lr=1e-3)\n",
    "ema_ebm1 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm1.load_state_dict(ebm1.state_dict())\n",
    "\n",
    "ebm2 = ToyModel().to(device)\n",
    "optimizer_ebm2 = optim.Adam(ebm2.parameters(), lr=1e-3)\n",
    "ema_ebm2 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm2.load_state_dict(ebm2.state_dict())\n",
    "\n",
    "ebm3 = ToyModel().to(device)\n",
    "optimizer_ebm3 = optim.Adam(ebm3.parameters(), lr=1e-3)\n",
    "ema_ebm3 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm3.load_state_dict(ebm3.state_dict())\n",
    "\n",
    "train(ebm1, ema_ebm1, optimizer_ebm1, dataset, compute_mse=False, algo='ed', m_particles=1, t_noise=1, w_stable=0.05)\n",
    "train(ebm2, ema_ebm2, optimizer_ebm2, dataset, compute_mse=False, algo='ed', m_particles=32, t_noise=1, w_stable=0.05)\n",
    "train(ebm3, ema_ebm3, optimizer_ebm3, dataset, compute_mse=False, algo='ed', m_particles=512, t_noise=1, w_stable=0.05)\n",
    "\n",
    "visualize_energynet(ema_ebm1, 'w=0.05, m=1')\n",
    "visualize_energynet(ema_ebm2, 'w=0.05, m=32')\n",
    "visualize_energynet(ema_ebm3, 'w=0.05, m=512')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$t=1, M=1, 32, 512, w=50$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50001/50001 [04:25<00:00, 188.62it/s]\n",
      "100%|██████████| 50001/50001 [04:38<00:00, 179.33it/s]\n",
      "100%|██████████| 50001/50001 [05:50<00:00, 142.76it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34236 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30988 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 25812 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm1 = ToyModel().to(device)\n",
    "optimizer_ebm1 = optim.Adam(ebm1.parameters(), lr=1e-3)\n",
    "ema_ebm1 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm1.load_state_dict(ebm1.state_dict())\n",
    "\n",
    "ebm2 = ToyModel().to(device)\n",
    "optimizer_ebm2 = optim.Adam(ebm2.parameters(), lr=1e-3)\n",
    "ema_ebm2 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm2.load_state_dict(ebm2.state_dict())\n",
    "\n",
    "ebm3 = ToyModel().to(device)\n",
    "optimizer_ebm3 = optim.Adam(ebm3.parameters(), lr=1e-3)\n",
    "ema_ebm3 = ToyModel(n_output=1).to(device)\n",
    "ema_ebm3.load_state_dict(ebm3.state_dict())\n",
    "\n",
    "train(ebm1, ema_ebm1, optimizer_ebm1, dataset, compute_mse=False, algo='ed', m_particles=1, t_noise=1, w_stable=50)\n",
    "train(ebm2, ema_ebm2, optimizer_ebm2, dataset, compute_mse=False, algo='ed', m_particles=32, t_noise=1, w_stable=50)\n",
    "train(ebm3, ema_ebm3, optimizer_ebm3, dataset, compute_mse=False, algo='ed', m_particles=512, t_noise=1, w_stable=50)\n",
    "\n",
    "visualize_energynet(ema_ebm1, 'w=50, m=1')\n",
    "visualize_energynet(ema_ebm2, 'w=50, m=32')\n",
    "visualize_energynet(ema_ebm3, 'w=50, m=512')"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
