hw10_flow_inpainting.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "638e3a9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAEICAYAAACZA4KlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAARYUlEQVR4nO3dfbBU9X3H8fcnxrRWiUIQpOj4VNtGS0Mq45goSMcGFcfR1OJoteJIC9aHaRxrVWKNjpPUOpqObUfaa3FEsT6hVIvxedKQDMYRGBAIUZESJdzwIBJx4gSEb//Yg13x7m8vu2cfLr/Pa+bO3Xu+e8757sJnzzl79uxPEYGZ7f0+0+kGzKw9HHazTDjsZplw2M0y4bCbZcJhN8uEwz5ASJou6T/Kvm8/lhWSfqeMZVlnyefZ20/SJcA1wNHA+8Bc4IaI2NLBtvokKYBjImJVH7X/AWZHRCkvLNZa3rK3maRrgH8ErgUOBE4EDgdekPS5GvN8tn0d2t7KYW8jSZ8HbgGuiohnI2J7RKwBzqMS+IuK+90saY6k2ZLeBy4pps2uWtbFkn4m6V1Jfy9pjaQ/qZp/dnH7iGJXfLKktyVtkvTNquWcIOllSVsk9Ur611ovOnUe23hJayX9naQNxbLOkTRR0huSNkua3t/1Spog6XVJv5R0t6QfSPrLqvqlklZKek/Sc5IO39Oec+Owt9dXgd8EnqieGBEfAM8AX6uafDYwBzgIeLD6/pKOBe4GLgRGUNlDGFln3ScDvwecCtwk6YvF9B3A1cBQ4CtF/fI9e1gfO4TK4xsJ3ATcQ+UF7HhgbLHeo+qtV9JQKo/9BuALwOtUnjuK+jnAdOBPgYOBHwIPNdhzNhz29hoKbIqIj/qo9Rb1XV6OiP+KiJ0R8eFu9/0z4L8j4kcRsY1KsOq9+XJLRHwYEUuBpcCXACJiUUT8OCI+KvYy/h04Zc8fGgDbgW9HxHbg4eLx3BURWyNiBbAC+MN+rHcisCIiniieq38GflG1nmnAP0TEyqL+HWC0t+5pDnt7bQKG1jgGH1HUd3knsZzfrq5HxK+Ad+usuzosvwIOAJD0u5LmSfpFccjwHT75orMn3o2IHcXtXS9Q66vqH/Zzvbs/vgDWVi3ncOCu4hBgC7AZEPX3brLmsLfXy8Cvqex+fkzS/sAZwEtVk1Nb6l7g0Kr596Oyu9uIGcBPqbzj/nkqu8dqcFllrXf3x6fqv6m8EEyLiIOqfvaLiAVt6HvActjbKCJ+SeUNun+RdLqkfSUdATxGZcv1QD8XNQc4S9JXize1bqHxgA6icvrvA0m/D/x1g8spc71PA6OKN/g+C1xB5f2AXf4NuEHScQCSDpQ0qU19D1gOe5tFxO1UtmJ3UPnP/gqVLdWpEfHrfi5jBXAVlePiXmArsIHKXsOe+lvgz4tl3AM80sAyGlFzvRGxCZgE3E7l8ORYYCHF44uIuVROXz5cHAIsp7JnZAn+UM1eQNIBwBYqu8T/2+F2SifpM1T2fC6MiO93up+Bylv2AUrSWZJ+qzjevwNYBqzpbFflkXSapIMk/Qb/fzz/4w63NaA57APX2cC64ucY4PzYu3bTvgK8ReUMxVnAOX2cgrQ94N14s0x4y26WibZeYFFcQWVmLRQRfZ6GbWrLXpwrfl3SKknXN7MsM2utho/ZJe0DvEHl4o21wKvABRHxk8Q83rKbtVgrtuwnAKsiYnVxMcbDVN4hNrMu1EzYR/LJizXW0seFCJKmSlooaWET6zKzJjXzBl1fuwqf2k2PiB6gB7wbb9ZJzWzZ1wKHVf19KJUPeJhZF2om7K8Cx0g6srjy6nzgqXLaMrOyNbwbHxEfSboSeA7YB7i3uBrLzLpQWz8u62N2s9ZryYdqzGzgcNjNMuGwm2XCYTfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhMNulomGh2y28gwbNixZf/TRR5P1BQsW1Kz19PQk512zZk2yvrc68MADk/Vx48Yl688++2yyvn379j3uqdWaCrukNcBWYAfwUUSMKaMpMytfGVv2P46ITSUsx8xayMfsZploNuwBPC9pkaSpfd1B0lRJCyUtbHJdZtaEZnfjT4qIdZKGAS9I+mlEzK++Q0T0AD0AkqLJ9ZlZg5raskfEuuL3BmAucEIZTZlZ+RoOu6T9JQ3adRuYACwvqzEzK1czu/HDgbmSdi3nPyMiffIxU4MHD07WV6xYkazXOye8fv36mrVcz6ND+nlbtGhRct6DDz44WR8zJn2W+c0330zWO6HhsEfEauBLJfZiZi3kU29mmXDYzTLhsJtlwmE3y4TDbpYJX+JagqFDhybrjzzySLI+ZMiQZP3uu+9O1q+66qpkPVc33nhjzdqRRx6ZnHfatGnJejeeWqvHW3azTDjsZplw2M0y4bCbZcJhN8uEw26WCYfdLBOKaN+Xx+yt31QzYcKEZP2ZZ55pavmHHHJIsr5x48amlj9QHXfcccn6smXLatbmzp2bnPeSSy5J1rdu3Zqsd1JEqK/p3rKbZcJhN8uEw26WCYfdLBMOu1kmHHazTDjsZpnw9ez9lBpW+dxzz21q2VOmTEnWfR69by+++GLDy653nr2bz6M3ylt2s0w47GaZcNjNMuGwm2XCYTfLhMNulgmH3SwTPs/eT3feeWfN2kUXXZSct97wwI899lhDPe3txo4dm6wPHz48Wb/vvvtq1mbPnt1ISwNa3S27pHslbZC0vGraEEkvSHqz+J0egNzMOq4/u/H3AafvNu164KWIOAZ4qfjbzLpY3bBHxHxg826TzwZmFbdnAeeU25aZla3RY/bhEdELEBG9kmp+cFzSVGBqg+sxs5K0/A26iOgBemDv/cJJs4Gg0VNv6yWNACh+byivJTNrhUbD/hQwubg9GXiynHbMrFXq7sZLeggYDwyVtBb4FnAb8KikKcDbwKRWNtkNUt+vv3PnzuS869atS9a3bdvWUE8DwX777VezNn369OS8l19+ebJeb8yDSy+9NFnPTd2wR8QFNUqnltyLmbWQPy5rlgmH3SwTDrtZJhx2s0w47GaZ8CWubXDmmWcm688//3yyvmXLlmR9xowZe9pSaU455ZRkffz48TVrJ554YlPrnjNnTlPz58ZbdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sE6p3mWCpKxvA31Rz/PHH16w9+WT6cv4RI0Y0tW5JyXo7/w1318reVq9enayffvru34P6SW+99VbD6x7IIqLPfxRv2c0y4bCbZcJhN8uEw26WCYfdLBMOu1kmHHazTPh69n5KDbs8atSo5LyjR49O1uudL7722muT9Y0bN9aszZo1q2atDA888ECyvnTp0oaXvWDBgmQ91/PojfKW3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhK9nt6YcddRRyfqqVatq1pYsWZKc97TTTkvWU58vyFnD17NLulfSBknLq6bdLOnnkpYUPxPLbNbMytef3fj7gL4+4vVPETG6+PleuW2ZWdnqhj0i5gOb29CLmbVQM2/QXSnptWI3f3CtO0maKmmhpIVNrMvMmtRo2GcARwOjgV7gzlp3jIieiBgTEWMaXJeZlaChsEfE+ojYERE7gXuAE8pty8zK1lDYJVV/N/LXgeW17mtm3aHu9eySHgLGA0MlrQW+BYyXNBoIYA0wrXUtWje76aabkvXU5ziuu+665Lw+j16uumGPiAv6mDyzBb2YWQv547JmmXDYzTLhsJtlwmE3y4TDbpYJf5W0JU2aNClZv/jii5P1rVu31qy9++67DfVkjfGW3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhM+zW9IZZ5zR1Pzz5s2rWVu8eHFTy7Y94y27WSYcdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJD9lsSb29vcn6oEGDkvVx48bVrPk8e2s0PGSzme0dHHazTDjsZplw2M0y4bCbZcJhN8uEw26Wif4M2XwYcD9wCLAT6ImIuyQNAR4BjqAybPN5EfFe61q1VrjsssuS9eHDhyfrGzZsSNZ9Lr179GfL/hFwTUR8ETgRuELSscD1wEsRcQzwUvG3mXWpumGPiN6IWFzc3gqsBEYCZwOzirvNAs5pUY9mVoI9OmaXdATwZeAVYHhE9ELlBQEYVnp3Zlaafn8HnaQDgMeBb0TE+1KfH7/ta76pwNTG2jOzsvRryy5pXypBfzAinigmr5c0oqiPAPp8pyYieiJiTESMKaNhM2tM3bCrsgmfCayMiO9WlZ4CJhe3JwNPlt+emZWlP7vxJwF/ASyTtKSYNh24DXhU0hTgbSA9tq91pXqn3updAv300083vO56l8cOHjw4WX/77bcbXneO6oY9In4E1DpAP7XcdsysVfwJOrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJD9lsTdmxY0eyfuGFF9asXX311cl5V6xYkaxPnjw5WbdP8pbdLBMOu1kmHHazTDjsZplw2M0y4bCbZcJhN8uEh2zO3JIlS5L1UaNGJev1vp4s9f9r5syZyXlvvfXWZP2dd95J1nPlIZvNMuewm2XCYTfLhMNulgmH3SwTDrtZJhx2s0z4PHvmxo4dm6zfcsstyfr8+fOT9RkzZtSsvfdeeoTvbdu2JevWN59nN8ucw26WCYfdLBMOu1kmHHazTDjsZplw2M0yUfc8u6TDgPuBQ4CdQE9E3CXpZuCvgI3FXadHxPfqLMvn2c1arNZ59v6EfQQwIiIWSxoELALOAc4DPoiIO/rbhMNu1nq1wl53RJiI6AV6i9tbJa0ERpbbnpm12h4ds0s6Avgy8Eox6UpJr0m6V9LgGvNMlbRQ0sLmWjWzZvT7s/GSDgB+AHw7Ip6QNBzYBARwK5Vd/UvrLMO78WYt1vAxO4CkfYF5wHMR8d0+6kcA8yLiD+osx2E3a7GGL4RR5etDZwIrq4NevHG3y9eB5c02aWat0593408Gfggso3LqDWA6cAEwmspu/BpgWvFmXmpZ3rKbtVhTu/FlcdjNWs/Xs5tlzmE3y4TDbpYJh90sEw67WSYcdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sEw67WSYcdrNM1P3CyZJtAn5W9ffQYlo36tbeurUvcG+NKrO3w2sV2no9+6dWLi2MiDEdayChW3vr1r7AvTWqXb15N94sEw67WSY6HfaeDq8/pVt769a+wL01qi29dfSY3czap9NbdjNrE4fdLBMdCbuk0yW9LmmVpOs70UMtktZIWiZpSafHpyvG0NsgaXnVtCGSXpD0ZvG7zzH2OtTbzZJ+Xjx3SyRN7FBvh0n6vqSVklZI+ptiekefu0RfbXne2n7MLmkf4A3ga8Ba4FXggoj4SVsbqUHSGmBMRHT8AxiSxgEfAPfvGlpL0u3A5oi4rXihHBwR13VJbzezh8N4t6i3WsOMX0IHn7syhz9vRCe27CcAqyJidURsAx4Gzu5AH10vIuYDm3ebfDYwq7g9i8p/lrar0VtXiIjeiFhc3N4K7BpmvKPPXaKvtuhE2EcC71T9vZbuGu89gOclLZI0tdPN9GH4rmG2it/DOtzP7uoO491Ouw0z3jXPXSPDnzerE2Hva2iabjr/d1JE/BFwBnBFsbtq/TMDOJrKGIC9wJ2dbKYYZvxx4BsR8X4ne6nWR19ted46Efa1wGFVfx8KrOtAH32KiHXF7w3AXCqHHd1k/a4RdIvfGzrcz8ciYn1E7IiIncA9dPC5K4YZfxx4MCKeKCZ3/Lnrq692PW+dCPurwDGSjpT0OeB84KkO9PEpkvYv3jhB0v7ABLpvKOqngMnF7cnAkx3s5RO6ZRjvWsOM0+HnruPDn0dE23+AiVTekX8L+GYneqjR11HA0uJnRad7Ax6islu3ncoe0RTgC8BLwJvF7yFd1NsDVIb2fo1KsEZ0qLeTqRwavgYsKX4mdvq5S/TVlufNH5c1y4Q/QWeWCYfdLBMOu1kmHHazTDjsZplw2M0y4bCbZeL/AKOmxh+z9/L6AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAEICAYAAACZA4KlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAARWUlEQVR4nO3dfbBU9X3H8fcnPsQRqYIPPImiju0kLYYoJXZUoBMThWYG7FQTSxVaG2BG22o6jlYZg1GbTKYx2klKioMCalCLEikmRms7pRlTRyBXBInhYVCQG0CEiE5TEL/9Y89Nl+vu2cs+nYXf5zWzc3fP7zx873I/nN+es+f8FBGY2ZHvY0UXYGbt4bCbJcJhN0uEw26WCIfdLBEOu1kiHHarm6TNki4tug7rG4e9hST9qaQVkt6T1C3pR5IuLrouAEnjJW1t4frnS7q7Veu3Q+ewt4ikrwD3AX8PDALOAP4JmFTHuo7uyzSzPA57C0g6EfgacH1EPBUR70fE/oj414i4OZvn45Luk7Qte9wn6eNZ23hJWyXdIumXwEOSZktaLOkRSe8C03rvPXvvrbNu9t9Jek3SbkkPSTpOUj/gR8DQrNfxnqShkj4m6VZJGyXtkvSEpIFl67tG0htZ2+2H8H6MkBSS/lzSlqyWmZJ+X9JqSXskfads/nMk/Xu2nbclPSrppLL28yX9TNJeSf8i6fFe78MXJHVl631R0nmH9A94hHLYW+MPgOOAJTnz3A5cCIwCPgWMAWaVtQ8GBgJnAtOzaZOAxcBJwKN9rGUKcBlwDvDbwKyIeB+YAGyLiBOyxzbgr4HJwDhgKLAb+C6ApE8Cc4BrsraTgdP7WEOPzwDnAl+k1Ou5HbgU+F3gKknjsvkEfD3bzieA4cDsrI5jKb2v8ym9P4uAK3o2IOl84EFgRlbjPwNLe/4jTVpE+NHkB6WA/bLGPBuBiWWvLwM2Z8/HA/uA48raZwPLe61jPnB32evxwNay15uBmWWvJwIbK82bTVsHfLbs9RBgP3A0cAfwWFlbv6zGS6v8fr+pDRgBBDCsrH0X8MWy108CN1ZZ12TgZ9nzscBbgMraf1K2rTnAXb2Wfx0YV/TfRdEPf+5rjV3AKZKOjogPqswzFHij7PUb2bQeOyPi172W2VJHLeXL9N5Gb2cCSyR9WDbtAKVjDkPL1xUR70vadYi1bC97/j8VXp8AIOk04B+BS4D+lHqgu7P5hgJvRZbiTPnveCYwVdJflU07lvzfOwnuxrfGT4FfU9ojVbON0h9mjzOyaT0qXY7Ye9r7wPFlrwdXWGZ4lW1UWv8WYEJEnFT2OC4i3gK6y9cl6XhK3eRW+HpW33kR8VvAn1Hq2pPVMUySyuYv/x23APf0+h2Oj4hFLar1sOGwt0BE/IpSt/e7kiZLOl7SMZImSPpmNtsiYJakUyWdks3/yCFuqguYKGmgpMHAjRXmuV7S6dmBttuAx7Pp24GTs4OJPb4H3CPpTICstp6zB4uBL0i6OPvc/DVa9/fTH3gP2CNpGHBzWdtPKfU2bpB0dFbfmLL2B4CZkj6jkn6S/khS/xbVethw2FskIu4FvkLpoNtOSnucG4AfZLPcDawAVgOvAquyaYfiYeAVSp/Nn+P/g1zu+1nbpuxxd1bfzyn9h7MpO2o9FLgfWAo8J2kv8N+UDqoREWuB67P1dVPqVrfqPP2dwPnAr4BngKd6GiJiH/DHwHXAHkp7/WXA/2btK4AvA9/JatwATGtRnYcVHfzRx44kkjYDfxkR/1Z0La0k6SXgexHxUNG1dDLv2e2wI2mcpMFZN34qcB7wbNF1dTofjbfD0e8AT1A6er8R+JOI6C62pM7nbrxZItyNN0tEW7vxktyNMGuxiFCl6Q3t2SVdLul1SRsk3drIusyster+zC7pKOAXwOconW99Gbg6Il7LWcZ7drMWa8WefQywISI2ZV90eIw6rtU2s/ZoJOzDOPgChK3ZtINImq7S3VpWNLAtM2tQIwfoKnUVPtJNj4i5wFxwN96sSI3s2bdy8NVGp3PwVVtm1kEaCfvLwLmSzsqugvoSpYsozKwD1d2Nj4gPJN0A/Bg4CngwuzLKzDpQW78u68/sZq3Xki/VmNnhw2E3S4TDbpYIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEQ67WSLqHrL5SNPO0WytPaSKg5kCcOKJJ+YuO3bs2Nz2Z599Nrd9//79ue1FaCjskjYDe4EDwAcRMboZRZlZ8zVjz/6HEfF2E9ZjZi3kz+xmiWg07AE8J2mlpOmVZpA0XdIKSSsa3JaZNaDRbvxFEbFN0mnA85J+HhHLy2eIiLnAXABJPgpmVpCG9uwRsS37uQNYAoxpRlFm1nx1h11SP0n9e54DnwfWNKswM2uuRrrxg4Al2bnMo4HvR0T+yUezNso7l75y5crcZU899dTc9tGj888yr1+/Pre9CHWHPSI2AZ9qYi1m1kI+9WaWCIfdLBEOu1kiHHazRDjsZonwJa52xJo1a1bVtrPOOit32RkzZuS2d+KptVq8ZzdLhMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEqF23kK5k+9U41tJH3ny/k2XLFmSu+y0adNy2/fu3VtPSW0RERXvoe09u1kiHHazRDjsZolw2M0S4bCbJcJhN0uEw26WCF/PbkmqdZ69k8+j18t7drNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sET7Pbkes+fPnV2175JFH2ldIh6i5Z5f0oKQdktaUTRso6XlJ67OfA1pbppk1qi/d+PnA5b2m3Qq8EBHnAi9kr82sg9UMe0QsB97pNXkSsCB7vgCY3NyyzKzZ6v3MPigiugEiolvSadVmlDQdmF7ndsysSVp+gC4i5gJzobNvOGl2pKv31Nt2SUMAsp87mleSmbVCvWFfCkzNnk8Fnm5OOWbWKjXvGy9pETAeOAXYDnwV+AHwBHAG8CZwZUT0PohXaV0d2433feOPPFLF26cf8ardN96DRGQc9iOPw34wf13WLBEOu1kiHHazRDjsZolw2M0S4UtcM7WO3C5cuLBq25QpUxra9vLly3Pb9+zZk9s+Z86chrbfiHHjxuW2jx8/vmrbhRde2NC2Fy9e3NDyqfGe3SwRDrtZIhx2s0Q47GaJcNjNEuGwmyXCYTdLhK9666MLLrigatvTT+dfzj9kyJCGtl3rOwBFXrHXyto2bdqU23755b3vg3qwjRs31r3tw5mvejNLnMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEuHr2fto5cqVVdtGjhyZu+yoUaNy22udL7755ptz23fu3Fm1bcGCBVXbmuHhhx/ObX/llVfqXveLL76Y257qefR6ec9ulgiH3SwRDrtZIhx2s0Q47GaJcNjNEuGwmyXC17NbQ84+++zc9g0bNlRt6+rqyl32sssuy23P+35Byuq+nl3Sg5J2SFpTNm22pLckdWWPic0s1syary/d+PlApa94fTsiRmWPHza3LDNrtpphj4jlwDttqMXMWqiRA3Q3SFqddfMHVJtJ0nRJKyStaGBbZtagesM+BzgHGAV0A9+qNmNEzI2I0RExus5tmVkT1BX2iNgeEQci4kPgAWBMc8sys2arK+ySyu+NfAWwptq8ZtYZal7PLmkRMB44RdJW4KvAeEmjgAA2AzNaV6J1sjvuuCO3Pe97HLfcckvusj6P3lw1wx4RV1eYPK8FtZhZC/nrsmaJcNjNEuGwmyXCYTdLhMNulgjfStpyXXnllbnt1157bW773r17q7bt2rWrrpqsPt6zmyXCYTdLhMNulgiH3SwRDrtZIhx2s0Q47GaJ8Hl2yzVhwoSGll+2bFnVtlWrVjW0bjs03rObJcJhN0uEw26WCIfdLBEOu1kiHHazRDjsZonwkM2Wq7u7O7e9f//+ue1jx46t2ubz7K1R95DNZnZkcNjNEuGwmyXCYTdLhMNulgiH3SwRDrtZIvoyZPNwYCEwGPgQmBsR90saCDwOjKA0bPNVEbG7daVaK8ycOTO3fdCgQbntO3bsyG33ufTO0Zc9+wfA30bEJ4ALgeslfRK4FXghIs4FXshem1mHqhn2iOiOiFXZ873AOmAYMAlYkM22AJjcohrNrAkO6TO7pBHAp4GXgEER0Q2l/xCA05penZk1TZ/vQSfpBOBJ4MaIeFeq+PXbSstNB6bXV56ZNUuf9uySjqEU9Ecj4qls8nZJQ7L2IUDFIzURMTciRkfE6GYUbGb1qRl2lXbh84B1EXFvWdNSYGr2fCrwdPPLM7Nm6Us3/iLgGuBVSV3ZtNuAbwBPSLoOeBPIH9vXOlKtU2+1LoF+5pln6t52rctjBwwYkNv+5ptv1r3tFNUMe0T8BKj2Af2zzS3HzFrF36AzS4TDbpYIh90sEQ67WSIcdrNEOOxmifCQzdaQAwcO5LZPmTKlattNN92Uu+zatWtz26dOnZrbbgfznt0sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4SHbE5cV1dXbvvIkSNz22vdnizv72vevHm5y95111257Vu2bMltT5WHbDZLnMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEuHz7Im75JJLctvvvPPO3Pbly5fnts+ZM6dq2+7d+SN879u3L7fdKvN5drPEOexmiXDYzRLhsJslwmE3S4TDbpYIh90sETXPs0saDiwEBgMfAnMj4n5Js4EvAzuzWW+LiB/WWJfPs5u1WLXz7H0J+xBgSESsktQfWAlMBq4C3ouIf+hrEQ67WetVC3vNEWEiohvozp7vlbQOGNbc8sys1Q7pM7ukEcCngZeySTdIWi3pQUkDqiwzXdIKSSsaK9XMGtHn78ZLOgH4T+CeiHhK0iDgbSCAuyh19f+ixjrcjTdrsbo/swNIOgZYBvw4Iu6t0D4CWBYRv1djPQ67WYvVfSGMSrcPnQesKw96duCuxxXAmkaLNLPW6cvR+IuB/wJepXTqDeA24GpgFKVu/GZgRnYwL29d3rObtVhD3fhmcdjNWs/Xs5slzmE3S4TDbpYIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEQ67WSIcdrNE1LzhZJO9DbxR9vqUbFon6tTaOrUucG31amZtZ1ZraOv17B/ZuLQiIkYXVkCOTq2tU+sC11avdtXmbrxZIhx2s0QUHfa5BW8/T6fW1ql1gWurV1tqK/Qzu5m1T9F7djNrE4fdLBGFhF3S5ZJel7RB0q1F1FCNpM2SXpXUVfT4dNkYejskrSmbNlDS85LWZz8rjrFXUG2zJb2VvXddkiYWVNtwSf8haZ2ktZL+Jpte6HuXU1db3re2f2aXdBTwC+BzwFbgZeDqiHitrYVUIWkzMDoiCv8ChqSxwHvAwp6htSR9E3gnIr6R/Uc5ICJu6ZDaZnOIw3i3qLZqw4xPo8D3rpnDn9ejiD37GGBDRGyKiH3AY8CkAuroeBGxHHin1+RJwILs+QJKfyxtV6W2jhAR3RGxKnu+F+gZZrzQ9y6nrrYoIuzDgC1lr7fSWeO9B/CcpJWSphddTAWDeobZyn6eVnA9vdUcxrudeg0z3jHvXT3DnzeqiLBXGpqmk87/XRQR5wMTgOuz7qr1zRzgHEpjAHYD3yqymGyY8SeBGyPi3SJrKVehrra8b0WEfSswvOz16cC2AuqoKCK2ZT93AEsofezoJNt7RtDNfu4ouJ7fiIjtEXEgIj4EHqDA9y4bZvxJ4NGIeCqbXPh7V6mudr1vRYT9ZeBcSWdJOhb4ErC0gDo+QlK/7MAJkvoBn6fzhqJeCkzNnk8Fni6wloN0yjDe1YYZp+D3rvDhzyOi7Q9gIqUj8huB24uooUpdZwOvZI+1RdcGLKLUrdtPqUd0HXAy8AKwPvs5sINqe5jS0N6rKQVrSEG1XUzpo+FqoCt7TCz6vcupqy3vm78ua5YIf4POLBEOu1kiHHazRDjsZolw2M0S4bCbJcJhN0vE/wHRKtDmgdzD7wAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAEICAYAAACZA4KlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAASv0lEQVR4nO3de6xdZZ3G8e9T7tAGWoG2ULQg1RlnxIodIlpLGUa5jAZIRKWidTC0GJsIcbgMZgSnccJMvIwhDEkNxVYUMEGgoQzCoFBKE6al0lKo0CKlFk7vLbQwYC+/+WOv45zWs9512JezTvs+n+Tk7L1/+137t3f7nLXWXnvtVxGBme3/BtXdgJn1D4fdLBMOu1kmHHazTDjsZplw2M0y4bDbgCDpi5IeqruP/Zl8nL1/SVoFDAd2AduBB4FpEbG9zr72JimAMRGxsgPLHg28BBwUETvbvXzrndfs9fhMRAwGxgIfBv6p3nbeOUkH1t2DvTMOe40iYi3wKxqhR9JHJS2QtFXSEkkTu+8raZik2yS9KmmLpHt71C6TtFLSZklzJB3XoxaSLpe0ohh3syQVtZMlPSbpNUkbJd1V3D6vGL5E0nZJn5c0UdIaSddIWgvcJukrkub3fE7F451cXD5M0vclvVw8xnxJhwHdy99aLP/0vZcl6WOSFhbjFkr6WI/ao5KmS3pC0jZJD0k6utV/j/1eRPinH3+AVcDfFZdHAc8APwKOBzYB59H4I/zJ4voxxX3nAncBQ4GDgDOK2/8W2AicChwC3ATM6/F4AdwPHAW8G9gAnFPU7gC+VTzeocD4vcad3OP6RGAn8G/F4xwGfAWYv9fz+9M44Gbg0eK5HQB8rBg7urjfgT3G/WlZwDBgC/Al4EDg4uL6u4r6o8CLwPuKPh4Fbqz733ag/3jNXo97JW0D/gCsB64HLgEeiIgHImJ3RDwMLALOkzQSOBe4PCK2RMSOiHisWNYXgZkRsTgi3qaxS3B6sV/c7caI2BoRq4HfUGxJADuA9wDHRcRbEbHHWroXu4HrI+LtiPjf1B0lDQIuBb4REa9ExK6IWFD0WOXvgRUR8dOI2BkRdwC/Az7T4z63RcQLRR+/6PGcrITDXo8LImIIjbXlXwBH0wjdRcUm/FZJW4HxwEjgBGBzRGzpZVnHAS93X4nGG32baKxNu63tcflNYHBx+WpAwP9IelbSpRV9b4iIt/r2FDmaxtbCi328f097PKfCy/TtOVkJh71Gxdr5J8D3aKzlfxoRR/X4OSIibixqwyQd1ctiXqXxhwIASUcA7wJe6cPjr42IyyLiOGAq8J/d+9tlQ/a6/gZweI/HHtGjthF4C3hvH5aztz2eU+Hd9OE5WTmHvX7/QWP/fD7wGUlnSzpA0qHFm2KjIqIL+C8aYRwq6SBJE4rxPwf+QdJYSYcA/wo8GRGrqh5Y0kWSRhVXt9AI4a7i+jrgpIpFLAH+qnjsQ4EbugsRsRuYCfxA0nHFczq96HEDjV2CsuU/ALxP0iRJB0r6PPABGu89WJMc9ppFxAZgNnAFcD5wHY0w/AG4iv//N/oSjX3s39HYz7+iGP8I8M/A3UAXjTXpF/r48H8DPClpOzCHxv71S0XtBmBWsUvxuZLeXwD+BfhvYAWNP1g9/SONNyAXAptpvLk3KCLeBL4LPFEs/6N7LXcT8GngmzR2Sa4GPh0RG/v4vKwX/lCNWSa8ZjfLhMNulgmH3SwTDrtZJvr1ZIbiTCoz66CIUG+3t7Rml3SOpOeLkzCubWVZZtZZTR96k3QA8AKND4SsoXEs9eKIeC4xxmt2sw7rxJr9NGBlRPw+Iv4I3EnjQyFmNgC1EvbjaXzKq9sa9jxRAQBJUyQtkrSohccysxa18gZdb5sKf7aZHhEzgBngzXizOrWyZl9D49TLbqNonK1kZgNQK2FfCIyRdKKkg2mcfDGnPW2ZWbs1vRkfETslTaPxHWoH0Pi2lGfb1pmZtVW/nvXmfXazzuvIh2rMbN/hsJtlwmE3y4TDbpYJh90sEw67WSYcdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sEw67WSYcdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sEw67WSYcdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sE01P2by/ufrqq5P1JUuWNFUDGDFiRLL+/PPPJ+tDhw5N1l9//fXS2vbt25NjL7zwwmT9nnvuSdaHDRuWrG/evLm0dtRRRyXHbt26NVk/5JBDkvW33367tHbkkUcmx06YMCFZf/DBB5P1HTt2JOt1aCnsklYB24BdwM6IGNeOpsys/dqxZj8zIja2YTlm1kHeZzfLRKthD+AhSU9JmtLbHSRNkbRI0qIWH8vMWtDqZvzHI+JVSccCD0v6XUTM63mHiJgBzACQFC0+npk1qaU1e0S8WvxeD9wDnNaOpsys/ZoOu6QjJA3pvgx8CljWrsbMrL0U0dyWtaSTaKzNobE78POI+G7FmH12M37w4MGlteHDhyfHjh8/Pllfvnx5sr5u3bpkPXUsfdOmTcmxrR7rrjJoUPn6ZPfu3cmxVcfRqxx66KGltaeeeio59phjjknWP/KRjyTrK1euTNY7KSLU2+1N77NHxO+BDzXdkZn1Kx96M8uEw26WCYfdLBMOu1kmHHazTPgU18KQIUOS9SuvvLK0dtZZZyXHfvCDH0zW58yZk6xfeumlyfqYMWNKa2+88UZy7IEHtvZfoJVTXKvcdNNNyXrVIcnUobfRo0cnx06dOjVZr/PQWrO8ZjfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuHj7IVt27Yl6yeeeGJp7f3vf39y7OrVq5P1BQsWJOtVnwGo+irqlLfeeitZP+yww5J1qdezKfvk9NNPT9Yff/zxZP22225L1lO93Xvvvcmxd911V7K+L/Ka3SwTDrtZJhx2s0w47GaZcNjNMuGwm2XCYTfLhI+zFw4//PBk/UMfKv8i3arpeavO+Z4/f36ynpqSuVWpr8iG1qd8njhxYmmt6hj9pEmTkvVWVE1FXfW890Ves5tlwmE3y4TDbpYJh90sEw67WSYcdrNMOOxmmfBx9sKbb76ZrO/cubO0VjX1cFdXV7I+YsSIZP3II49M1teuXVtaqzpffdSoUcl61XH0qu9unzt3bmntsssuS45duHBhsn7GGWck6zNnziyt3X777cmx+6PKNbukmZLWS1rW47Zhkh6WtKL4PbSzbZpZq/qyGf8T4Jy9brsWeCQixgCPFNfNbACrDHtEzAP2nsPnfGBWcXkWcEF72zKzdmt2n314RHQBRESXpGPL7ihpCjClyccxszbp+Bt0ETEDmAEgKTr9eGbWu2YPva2TNBKg+L2+fS2ZWSc0G/Y5wOTi8mTgvva0Y2adooj0lrWkO4CJwNHAOuB64F7gF8C7gdXARRFRORH3vrwZ/9JLL5XWqs6Fr3LnnXcm66+99lqynvre+Kq54avOlX/iiSeS9enTpyfry5YtK6197WtfS46tMm3atGT95ptvbmn5+6qI6PWLAir32SPi4pLSWS11ZGb9yh+XNcuEw26WCYfdLBMOu1kmHHazTPgU1z6aN29eae2SSy5Jjq06vHnKKack61u3bk3WU19F/etf/zo5tqq3gw8+OFlPvS4AZ555Zmmt6tTgKhMmTEjWcz30VsZrdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sE5WnuLb1wfbhU1xPPfXU0tp996VP5x85cmSyXjV1cVW9k/+GVcseNCi9vkiNr1p26rRigHPO2ft7UPf04osvJuv7q7JTXL1mN8uEw26WCYfdLBMOu1kmHHazTDjsZplw2M0y4fPZ+2jx4sWltaqvax47dmyyfu655ybrV111VbK+YcOG0tqsWbNKa+0we/bsZH3p0qWltarPDyxYsCBZz/U4erO8ZjfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuHz2a0lJ510UrK+YsWK0tqSJUuSY88+++xkPfX5gpw1fT67pJmS1kta1uO2GyS9Iunp4ue8djZrZu3Xl834nwC9fSXIDyNibPHzQHvbMrN2qwx7RMwDNvdDL2bWQa28QTdN0tJiM39o2Z0kTZG0SNKiFh7LzFrUbNhvAd4LjAW6gO+X3TEiZkTEuIgY1+RjmVkbNBX2iFgXEbsiYjfwY+C09rZlZu3WVNgl9fxu5AuBZWX3NbOBofJ8dkl3ABOBoyWtAa4HJkoaCwSwCpjauRZtIPv2t7/d9NhrrrkmWfdx9PaqDHtEXNzLzbd2oBcz6yB/XNYsEw67WSYcdrNMOOxmmXDYzTLhr5K2pM9+9rPJ+pe//OVkfdu2baW1TZs2NdWTNcdrdrNMOOxmmXDYzTLhsJtlwmE3y4TDbpYJh90sEz7ObklV00lXuf/++0trqWmwrf28ZjfLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMuEpmy2pq6srWR8yZEiyPmHChNKaj7N3RtNTNpvZ/sFhN8uEw26WCYfdLBMOu1kmHHazTDjsZpnoy5TNJwCzgRHAbmBGRPxI0jDgLmA0jWmbPxcRWzrXqnXC1Knp2baPPfbYZL1qWmUfSx84+rJm3wl8MyL+Evgo8HVJHwCuBR6JiDHAI8V1MxugKsMeEV0Rsbi4vA1YDhwPnA/MKu42C7igQz2aWRu8o312SaOBDwNPAsMjogsafxCA9PaemdWqz99BJ2kwcDdwRUS8LvX68dvexk0BpjTXnpm1S5/W7JIOohH0n0XEL4ub10kaWdRHAut7GxsRMyJiXESMa0fDZtacyrCrsQq/FVgeET/oUZoDTC4uTwbua397ZtYufdmM/zjwJeAZSU8Xt10H3Aj8QtJXgdXARR3p0Drq8ssvb2n83Llzmx5bdXrs0KFDk/XVq1c3/dg5qgx7RMwHynbQz2pvO2bWKf4EnVkmHHazTDjsZplw2M0y4bCbZcJhN8uEp2y2luzatStZnzRpUmntyiuvTI597rnnkvXJkycn67Ynr9nNMuGwm2XCYTfLhMNulgmH3SwTDrtZJhx2s0x4yubM/fa3v03WTznllJaWn/r/deuttybHTp8+PVlfs2ZNUz3t7zxls1nmHHazTDjsZplw2M0y4bCbZcJhN8uEw26WCR9nz9wnPvGJZP073/lOsv7YY48l67fccktpbcuW9AzfO3bsSNatdz7ObpY5h90sEw67WSYcdrNMOOxmmXDYzTLhsJtlovI4u6QTgNnACGA3MCMifiTpBuAyYENx1+si4oGKZfk4u1mHlR1n70vYRwIjI2KxpCHAU8AFwOeA7RHxvb424bCbdV5Z2CtnhImILqCruLxN0nLg+Pa2Z2ad9o722SWNBj4MPFncNE3SUkkzJQ0tGTNF0iJJi1pr1cxa0efPxksaDDwGfDcifilpOLARCGA6jU39SyuW4c14sw5rep8dQNJBwP3AryLiB73URwP3R8RfVyzHYTfrsKZPhJEk4FZgec+gF2/cdbsQWNZqk2bWOX15N3488DjwDI1DbwDXARcDY2lsxq8CphZv5qWW5TW7WYe1tBnfLg67Wef5fHazzDnsZplw2M0y4bCbZcJhN8uEw26WCYfdLBMOu1kmHHazTDjsZplw2M0y4bCbZcJhN8uEw26WicovnGyzjcDLPa4fXdw2EA3U3gZqX+DemtXO3t5TVujX89n/7MGlRRExrrYGEgZqbwO1L3Bvzeqv3rwZb5YJh90sE3WHfUbNj58yUHsbqH2Be2tWv/RW6z67mfWfutfsZtZPHHazTNQSdknnSHpe0kpJ19bRQxlJqyQ9I+npuuenK+bQWy9pWY/bhkl6WNKK4nevc+zV1NsNkl4pXrunJZ1XU28nSPqNpOWSnpX0jeL2Wl+7RF/98rr1+z67pAOAF4BPAmuAhcDFEfFcvzZSQtIqYFxE1P4BDEkTgO3A7O6ptST9O7A5Im4s/lAOjYhrBkhvN/AOp/HuUG9l04x/hRpfu3ZOf96MOtbspwErI+L3EfFH4E7g/Br6GPAiYh6wea+bzwdmFZdn0fjP0u9KehsQIqIrIhYXl7cB3dOM1/raJfrqF3WE/XjgDz2ur2FgzfcewEOSnpI0pe5mejG8e5qt4vexNfezt8ppvPvTXtOMD5jXrpnpz1tVR9h7m5pmIB3/+3hEnAqcC3y92Fy1vrkFeC+NOQC7gO/X2UwxzfjdwBUR8XqdvfTUS1/98rrVEfY1wAk9ro8CXq2hj15FxKvF7/XAPTR2OwaSdd0z6Ba/19fcz59ExLqI2BURu4EfU+NrV0wzfjfws4j4ZXFz7a9db3311+tWR9gXAmMknSjpYOALwJwa+vgzko4o3jhB0hHApxh4U1HPASYXlycD99XYyx4GyjTeZdOMU/NrV/v05xHR7z/AeTTekX8R+FYdPZT0dRKwpPh5tu7egDtobNbtoLFF9FXgXcAjwIri97AB1NtPaUztvZRGsEbW1Nt4GruGS4Gni5/z6n7tEn31y+vmj8uaZcKfoDPLhMNulgmH3SwTDrtZJhx2s0w47GaZcNjNMvF/d35Vd29egOQAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image, make_grid\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(action='ignore')\n",
    "\n",
    "batch_size = 128\n",
    "(full_dim, mid_dim, hidden) = (1 * 28 * 28, 1000, 5)\n",
    "lr = 1e-3\n",
    "epochs = 100\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "############################################\n",
    "# STEP 1: Define dataset and preprocessing #\n",
    "############################################\n",
    "\n",
    "\n",
    "class Logistic(torch.distributions.Distribution):\n",
    "    def __init__(self):\n",
    "        super(Logistic, self).__init__()\n",
    "\n",
    "    def log_prob(self, x):\n",
    "        return -(F.softplus(x) + F.softplus(-x))\n",
    "\n",
    "    def sample(self, size):\n",
    "        z = torch.distributions.Uniform(0., 1.).sample(size).to(device)\n",
    "        return torch.log(z) - torch.log(1. - z)\n",
    "    \n",
    "####################################\n",
    "# STEP 3: Implement Coupling Layer #\n",
    "####################################\n",
    "\n",
    "class Coupling(nn.Module):\n",
    "    def __init__(self, in_out_dim, mid_dim, hidden, mask_config):\n",
    "        super(Coupling, self).__init__()\n",
    "        self.mask_config = mask_config\n",
    "        \n",
    "        self.in_block = nn.Sequential(nn.Linear(in_out_dim//2, mid_dim), nn.ReLU())\n",
    "        self.mid_block = nn.ModuleList([nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU())\n",
    "                                                                 for _ in range(hidden - 1)])\n",
    "        self.out_block = nn.Linear(mid_dim, in_out_dim//2)\n",
    "\n",
    "    def forward(self, x, reverse=False):\n",
    "        [B, W] = list(x.size())\n",
    "        x = x.reshape((B, W//2, 2))\n",
    "        if self.mask_config:\n",
    "            on, off = x[:, :, 0], x[:, :, 1]\n",
    "        else:\n",
    "            off, on = x[:, :, 0], x[:, :, 1]\n",
    "\n",
    "        off_ = self.in_block(off)\n",
    "        for i in range(len(self.mid_block)):\n",
    "            off_ = self.mid_block[i](off_)\n",
    "        shift = self.out_block(off_)\n",
    "        \n",
    "        if reverse:\n",
    "            on = on - shift\n",
    "        else:\n",
    "            on = on + shift\n",
    "\n",
    "        if self.mask_config:\n",
    "            x = torch.stack((on, off), dim=2)\n",
    "        else:\n",
    "            x = torch.stack((off, on), dim=2)\n",
    "        return x.reshape((B, W))\n",
    "\n",
    "class Scaling(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super(Scaling, self).__init__()\n",
    "        self.scale = nn.Parameter(torch.zeros((1, dim)), requires_grad=True)\n",
    "\n",
    "    def forward(self, x, reverse=False):\n",
    "        log_det_J = torch.sum(self.scale)\n",
    "        if reverse:\n",
    "            x = x * torch.exp(-self.scale)\n",
    "        else:\n",
    "            x = x * torch.exp(self.scale)\n",
    "        return x, log_det_J\n",
    "    \n",
    "##########################\n",
    "# STEP 4: Implement NICE #\n",
    "##########################\n",
    "\n",
    "class NICE(nn.Module):\n",
    "    def __init__(self,in_out_dim, mid_dim, hidden, mask_config=1.0, coupling=4):\n",
    "        super(NICE, self).__init__()\n",
    "        self.prior = Logistic()\n",
    "        self.in_out_dim = in_out_dim\n",
    "\n",
    "        self.coupling = nn.ModuleList([\n",
    "            Coupling(in_out_dim=in_out_dim, \n",
    "                     mid_dim=mid_dim, \n",
    "                     hidden=hidden, \n",
    "                     mask_config=(mask_config+i)%2) \\\n",
    "            for i in range(coupling)])\n",
    "        \n",
    "        self.scaling = Scaling(in_out_dim)\n",
    "\n",
    "    def g(self, z):\n",
    "        x, _ = self.scaling(z, reverse=True)\n",
    "        for i in reversed(range(len(self.coupling))):\n",
    "            x = self.coupling[i](x, reverse=True)\n",
    "        return x\n",
    "\n",
    "    def f(self, x):\n",
    "        for i in range(len(self.coupling)):\n",
    "            x = self.coupling[i](x)\n",
    "        z, log_det_J = self.scaling(x)\n",
    "        return z, log_det_J\n",
    "\n",
    "    def log_prob(self, x):\n",
    "        z, log_det_J = self.f(x)\n",
    "        log_ll = torch.sum(self.prior.log_prob(z), dim=1)\n",
    "        return log_ll + log_det_J\n",
    "\n",
    "    def sample(self, size):\n",
    "        z = self.prior.sample((size, self.in_out_dim)).to(device)\n",
    "        return self.g(z)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.log_prob(x)\n",
    "\n",
    "\n",
    "# Load pre-trained NICE model onto CPU\n",
    "model = NICE(in_out_dim=784, mid_dim=1000, hidden=5).to(device)\n",
    "model.load_state_dict(torch.load('nice.pt',map_location=torch.device('cpu')))\n",
    "\n",
    "# Since we do not update model, set requires_grad = False\n",
    "model.requires_grad_(False)\n",
    "\n",
    "# Get an MNIST image\n",
    "testset = torchvision.datasets.MNIST(root='./', train=False, download=True, transform=torchvision.transforms.ToTensor())\n",
    "test_loader = DataLoader(dataset = testset, batch_size=1, shuffle=False)\n",
    "pass_count = 6\n",
    "itr = iter(test_loader)\n",
    "for _ in range(pass_count+1):\n",
    "    image, _ = next(itr)\n",
    "    \n",
    "plt.figure(figsize = (4,4))\n",
    "plt.title('Original Image')\n",
    "plt.imshow(make_grid(image.squeeze().detach()).permute(1,2,0))\n",
    "#plt.show()\n",
    "plt.savefig('plt1.png')\n",
    "\n",
    "\n",
    "# Create mask \n",
    "mask = torch.ones_like(image,dtype=torch.bool)\n",
    "mask[:,:,5:12,5:20] = 0\n",
    "\n",
    "# Partially corrupt the image\n",
    "image[mask.logical_not()] = torch.ones_like(image[mask.logical_not()])\n",
    "plt.figure(figsize = (4,4))\n",
    "plt.title('Corrupted Image')\n",
    "plt.imshow(make_grid(image.squeeze()).permute(1,2,0))\n",
    "#plt.show()\n",
    "plt.savefig('plt2.png')\n",
    "\n",
    "\n",
    "lr = 1e-3\n",
    "X = image.clone().requires_grad_(True)\n",
    "optimizer = optim.SGD([X], lr=lr)\n",
    "\n",
    "epochs = 300\n",
    "\n",
    "for i in range(epochs):\n",
    "    # print(\"{}th epoch starting.\".format(i))\n",
    "    X.requires_grad_(True)\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    loss = - model(X.view(-1,784)).mean()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        torch.clamp(X, min=0, max=1)\n",
    "        X.data[mask] = image[mask]\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    Z, _ = model.f(X.view(-1,784))\n",
    "    recon = model.g(Z)\n",
    "    recon = recon.view(-1,1,28,28)\n",
    "\n",
    "# Plot reconstruction\n",
    "plt.figure(figsize = (4,4))\n",
    "plt.title('Reconstruction')\n",
    "plt.imshow(make_grid(recon.squeeze().detach()).permute(1,2,0))\n",
    "plt.show()\n",
    "plt.savefig('plt3.png')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519112a6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}