{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 torch.nn as nn\n",
    "torch.cuda.set_device(0)\n",
    "\n",
    "nz = 64\n",
    "t_sqrt = 0.5\n",
    "m_particles = 16\n",
    "batch_size = 256\n",
    "\n",
    "e_init_sig = 1.0\n",
    "e_prior_sig = 1.0\n",
    "e_l_steps = 500\n",
    "e_l_step_size = 0.2\n",
    "\n",
    "g_l_steps = 20\n",
    "g_l_step_size = 0.1\n",
    "g_llhd_sigma = 0.3\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        ngf = 128\n",
    "        self.gen = nn.Sequential(\n",
    "            nn.ConvTranspose2d(nz, ngf*8, 4, 1, 0, bias=True),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.ConvTranspose2d(ngf*8, ngf*4, 3, 2, 1, bias=True),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=True),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.ConvTranspose2d(ngf*2, 1, 4, 2, 1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "    def forward(self, z):\n",
    "        return self.gen(z)\n",
    "\n",
    "class EBMPrior(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        ndf = 200\n",
    "        self.ebm = nn.Sequential(\n",
    "            nn.Linear(nz, ndf),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Linear(ndf, ndf),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Linear(ndf, ndf),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Linear(ndf, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, z):\n",
    "        return self.ebm(z.squeeze())\n",
    "    \n",
    "def weights_init_xavier(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        nn.init.xavier_normal_(m.weight)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        m.weight.data.normal_(1, 0.02)\n",
    "        m.bias.data.fill_(0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sampling Algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mse = nn.MSELoss(reduction='sum')\n",
    "\n",
    "def sample_p_0(n=batch_size, sig=e_init_sig):\n",
    "    return sig * torch.randn(*[n, nz, 1, 1]).to(device)\n",
    "\n",
    "def sample_langevin_prior_z(z, netE):\n",
    "    z = z.clone().detach()\n",
    "    z.requires_grad = True\n",
    "\n",
    "    z_list = [z.clone().detach()]\n",
    "    for _ in range(e_l_steps):\n",
    "        en = netE(z)\n",
    "        z_grad = torch.autograd.grad(en.sum(), z)[0]\n",
    "        z.data = z.data - 0.5 * e_l_step_size * e_l_step_size * (z_grad + 1.0 / (e_prior_sig * e_prior_sig) * z.data)\n",
    "        z.data += e_l_step_size * torch.randn_like(z).data\n",
    "        z_list.append(z.clone().detach())\n",
    "    \n",
    "    return z_list\n",
    "\n",
    "def sample_langevin_post_z(z, x, netG, netE):\n",
    "    z = z.clone().detach()\n",
    "    z.requires_grad = True\n",
    "    for i in range(g_l_steps):\n",
    "        x_hat = netG(z)\n",
    "        g_log_lkhd = 1.0 / (2.0 * g_llhd_sigma * g_llhd_sigma) * mse(x_hat, x)\n",
    "        z_grad_g = torch.autograd.grad(g_log_lkhd, z)[0]\n",
    "\n",
    "        en = netE(z)\n",
    "        z_grad_e = torch.autograd.grad(en.sum(), z)[0]\n",
    "\n",
    "        z.data = z.data - 0.5 * g_l_step_size * g_l_step_size * (z_grad_g + z_grad_e + 1.0 / (e_prior_sig * e_prior_sig) * z.data)\n",
    "        z.data += g_l_step_size * torch.randn_like(z).data\n",
    "\n",
    "    return z.detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Eoch:    0, loss_g: 42.6653, loss_e: 0.2657: 100%|██████████| 235/235 [06:41<00:00,  1.71s/it]\n",
      "Eoch:    1, loss_g: 31.3591, loss_e: 0.0630: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    2, loss_g: 27.2527, loss_e: 0.0787: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    3, loss_g: 24.4436, loss_e: 0.0447: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    4, loss_g: 21.5762, loss_e: 0.2943: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    5, loss_g: 20.8515, loss_e: 0.1253: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    6, loss_g: 24.2443, loss_e: 0.1080: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    7, loss_g: 25.0665, loss_e: 0.1617: 100%|██████████| 235/235 [06:46<00:00,  1.73s/it]\n",
      "Eoch:    8, loss_g: 18.1657, loss_e: 0.2102: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:    9, loss_g: 20.5421, loss_e: 0.0733: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   10, loss_g: 21.7274, loss_e: 0.2285: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   11, loss_g: 17.1042, loss_e: 0.1408: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   12, loss_g: 19.4898, loss_e: 0.1105: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   13, loss_g: 18.1517, loss_e: 0.0912: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   14, loss_g: 17.1051, loss_e: 0.1519: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   15, loss_g: 16.7297, loss_e: 0.0741: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   16, loss_g: 19.0070, loss_e: 0.0881: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   17, loss_g: 14.9170, loss_e: 0.1606: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   18, loss_g: 15.0022, loss_e: 0.2015: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   19, loss_g: 17.3658, loss_e: 0.1596: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   20, loss_g: 16.1203, loss_e: 0.1186: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   21, loss_g: 16.4657, loss_e: 0.1792: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   22, loss_g: 18.1705, loss_e: 0.2099: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   23, loss_g: 15.6138, loss_e: 0.0858: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   24, loss_g: 15.6145, loss_e: 0.1996: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   25, loss_g: 17.0199, loss_e: 0.1594: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   26, loss_g: 16.9669, loss_e: 0.1827: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   27, loss_g: 18.3790, loss_e: 0.2206: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   28, loss_g: 17.5240, loss_e: 0.1693: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n",
      "Eoch:   29, loss_g: 16.3309, loss_e: 0.0739: 100%|██████████| 235/235 [06:45<00:00,  1.73s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "transform = transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                ])\n",
    "ds_train = datasets.MNIST('./datasets', train=True, download=False, transform=transform)\n",
    "dataloader_train = torch.utils.data.DataLoader(ds_train, batch_size=batch_size, shuffle=True, num_workers=0)\n",
    "\n",
    "netG = Generator().apply(weights_init_xavier).to(device)\n",
    "netE = EBMPrior().apply(weights_init_xavier).to(device)\n",
    "ema_netE = EBMPrior().to(device)\n",
    "ema_netE.load_state_dict(netE.state_dict())\n",
    "\n",
    "optG = torch.optim.Adam(netG.parameters(), lr=0.0001, betas=(0.5, 0.999))\n",
    "optE = torch.optim.Adam(netE.parameters(), lr=0.00005, betas=(0.5, 0.999))\n",
    "\n",
    "lr_scheduleE = torch.optim.lr_scheduler.ExponentialLR(optE, 0.99)\n",
    "lr_scheduleG = torch.optim.lr_scheduler.ExponentialLR(optG, 0.99)\n",
    "\n",
    "def eval_flag():\n",
    "    netG.eval()\n",
    "    netE.eval()\n",
    "\n",
    "def train_flag():\n",
    "    netG.train()\n",
    "    netE.train()\n",
    "\n",
    "for epoch in range(30):\n",
    "    pbar = tqdm(dataloader_train)\n",
    "    for x, _ in pbar:\n",
    "        x = x.to(device) * 2 - 1\n",
    "\n",
    "        z_g_0 = sample_p_0(x.size(0))\n",
    "        z_g_k = sample_langevin_post_z(Variable(z_g_0), x, netG, netE)\n",
    "\n",
    "        # Learn generator\n",
    "        optG.zero_grad()\n",
    "        x_hat = netG(z_g_k.detach())\n",
    "        loss_g = mse(x_hat, x) / x.size(0)\n",
    "        loss_g.backward()\n",
    "        optG.step()\n",
    "\n",
    "        # Learn prior EBM\n",
    "        optE.zero_grad()\n",
    "        z_g_k = z_g_k.detach().squeeze()    # [bs, z_dim]\n",
    "        pert_z_g_k = z_g_k + t_sqrt * torch.randn_like(z_g_k)\n",
    "        pert_z_g_k = pert_z_g_k.unsqueeze(1) + t_sqrt * \\\n",
    "                torch.randn(z_g_k.size(0), m_particles, *z_g_k.shape[1:]).to(z_g_k.device) # [bs, m_particles, z_dim]\n",
    "        pert_z_g_k = pert_z_g_k.detach().view(-1, *pert_z_g_k.shape[2:])  # [bs * m_particles, z_dim]\n",
    "        \n",
    "        en_pos = netE(z_g_k).view(z_g_k.size(0)) \\\n",
    "                    + 1.0 / (2 * e_prior_sig * e_prior_sig) * (z_g_k**2).sum(dim=-1)  # [bs]\n",
    "        en_neg = netE(pert_z_g_k).view(z_g_k.size(0), -1) \\\n",
    "                    + 1.0 / (2 * e_prior_sig * e_prior_sig) * (pert_z_g_k**2).sum(dim=-1).view(z_g_k.size(0), -1)   # [bs, m_particles]\n",
    "        en_ctr = en_pos.unsqueeze(-1) - en_neg  # [bs, m_particles]\n",
    "        en_ctr = torch.cat([torch.zeros_like(en_ctr[:, :1]), en_ctr], dim=1)    # [bs, m_particles + 1]\n",
    "        loss_e = en_ctr.logsumexp(dim=-1).mean()    # [bs]\n",
    "        loss_e.backward()\n",
    "        optE.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for p, ema_p in zip(netE.parameters(), ema_netE.parameters()):\n",
    "                ema_p.mul_(0.999).add_(p, alpha=0.001)\n",
    "\n",
    "        pbar.set_description('Eoch: {:4d}, loss_g: {:.4f}, loss_e: {:.4f}'.format(epoch, loss_g.item(), loss_e.item()))\n",
    "    \n",
    "    lr_scheduleE.step()\n",
    "    lr_scheduleG.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63836 characters]",
      "text/plain": [
       "<Figure size 300x300 with 64 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 72304 characters]",
      "text/plain": [
       "<Figure size 300x300 with 64 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "eval_flag()\n",
    "\n",
    "x = next(iter(dataloader_train))[0][:64, ...].to(device)\n",
    "z_g_0 = sample_p_0(n=x.size(0))\n",
    "z_g_k = sample_langevin_post_z(Variable(z_g_0), x*2-1, netG, netE)\n",
    "\n",
    "fig = plt.figure(figsize=(3, 3), constrained_layout=True)\n",
    "gs = fig.add_gridspec(8, 8)\n",
    "imgs = x.detach().cpu().reshape(8, 8, 28, 28).numpy()\n",
    "for n_row in range(8):\n",
    "    for n_col in range(8):\n",
    "        f_ax = fig.add_subplot(gs[n_row, n_col])\n",
    "        f_ax.imshow((imgs[n_row, n_col]+1.0) * 255 / 2, cmap=\"gray\")\n",
    "        f_ax.axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure(figsize=(3, 3), constrained_layout=True)\n",
    "gs = fig.add_gridspec(8, 8)\n",
    "imgs = netG(z_g_k).detach().cpu().reshape(8, 8, 28, 28).numpy()\n",
    "for n_row in range(8):\n",
    "    for n_col in range(8):\n",
    "        f_ax = fig.add_subplot(gs[n_row, n_col])\n",
    "        f_ax.imshow((imgs[n_row, n_col]+1.0) * 255 / 2, cmap=\"gray\")\n",
    "        f_ax.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 75968 characters]",
      "text/plain": [
       "<Figure size 300x300 with 64 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "z_e_0 = sample_p_0(n=64)\n",
    "z_e_ks = sample_langevin_prior_z(Variable(z_e_0), netE)\n",
    "\n",
    "fig = plt.figure(figsize=(3, 3), constrained_layout=True)\n",
    "gs = fig.add_gridspec(8, 8)\n",
    "imgs = netG(z_e_ks[-1]).detach().cpu().reshape(8, 8, 28, 28).numpy()\n",
    "for n_row in range(8):\n",
    "    for n_col in range(8):\n",
    "        f_ax = fig.add_subplot(gs[n_row, n_col])\n",
    "        f_ax.imshow((imgs[n_row, n_col]+1.0) * 255 / 2, cmap=\"gray\")\n",
    "        f_ax.axis(\"off\")\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
