{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prerequisites\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "from torchvision.utils import save_image\n",
    "from utils import make_new_MNIST_datasets, transform, lightweight_denoising_loss, compose2\n",
    "import torch.utils.data as data_utils\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# VAE and training from https://github.com/lyeoni/pytorch-mnist-VAE\n",
    "\n",
    "bs = 30\n",
    "P = bs\n",
    "# MNIST Dataset\n",
    "MNIST_DIR = '/mnt/home/hlawrence/ceph/datasets' #'some/path/to/dataset'\n",
    "train_dataset = datasets.MNIST(root=MNIST_DIR, train=True, transform=transforms.ToTensor(), download=True)\n",
    "test_dataset = datasets.MNIST(root='/mnt/home/hlawrence/ceph/datasets', train=False, transform=transforms.ToTensor(), download=False)\n",
    "\n",
    "# Data Loader (Input Pipeline)\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=bs, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=bs, shuffle=False)\n",
    "\n",
    "MNIST_flip = False\n",
    "MNIST_blur = True\n",
    "MNIST_color = False\n",
    "ds, dl, dl_iterator, ds_test, dl_test, dl_test_iterator = make_new_MNIST_datasets(lambda im: \n",
    "                                                                                  transform(im, flip=MNIST_flip, \n",
    "                                                                                            blur=MNIST_blur, \n",
    "                                                                                            colorjitter=MNIST_color), bs, MNIST_DIR=MNIST_DIR)\n",
    "\n",
    "def make_test_sample_functions(dl_iterator=dl_iterator, dl_test_iterator=dl_test_iterator):\n",
    "    def test_sample_func_train(dl_iterator=dl_iterator):\n",
    "        with torch.no_grad():\n",
    "            temp, _ = next(dl_iterator)\n",
    "            return temp.view(temp.shape[0], -1).float().to(device)\n",
    "    def test_sample_func_test(dl_test_iterator=dl_test_iterator):\n",
    "        with torch.no_grad():\n",
    "            temp, _ = next(dl_test_iterator)\n",
    "            return temp.view(temp.shape[0], -1).float().to(device)\n",
    "    return test_sample_func_train, test_sample_func_test\n",
    "test_sample_func_train, test_sample_func_test = make_test_sample_functions(dl_iterator=iter(dl), dl_test_iterator=iter(dl_test))\n",
    "\n",
    "totalnum = 50\n",
    "indices = torch.arange(totalnum)\n",
    "tr_10k = data_utils.Subset(ds, indices)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_10k\n",
    "dl_10k = torch.utils.data.DataLoader(dataset=tr_10k, batch_size=bs, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(784, 400)\n",
    "        self.fc21 = nn.Linear(400, 20)\n",
    "        self.fc22 = nn.Linear(400, 20)\n",
    "        self.fc3 = nn.Linear(20, 400)\n",
    "        self.fc4 = nn.Linear(400, 784)\n",
    "\n",
    "    def encode(self, x):\n",
    "        h1 = F.relu(self.fc1(x))\n",
    "        return self.fc21(h1), self.fc22(h1)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5*logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "\n",
    "    def decode(self, z):\n",
    "        h3 = F.relu(self.fc3(z))\n",
    "        return torch.sigmoid(self.fc4(h3))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x.view(-1, 784))\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar\n",
    "vae = VAE()\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "    vae.cuda()\n",
    "else:\n",
    "    device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====> Epoch: 10 Average loss: 2.7481\n",
      "====> Epoch: 20 Average loss: 2.7468\n",
      "====> Epoch: 30 Average loss: 2.7450\n",
      "====> Epoch: 40 Average loss: 2.7467\n",
      "====> Epoch: 50 Average loss: 2.7471\n",
      "====> Epoch: 60 Average loss: 2.7478\n",
      "====> Epoch: 70 Average loss: 2.7471\n",
      "====> Epoch: 80 Average loss: 2.7467\n",
      "====> Epoch: 90 Average loss: 2.7465\n",
      "====> Epoch: 100 Average loss: 2.7454\n",
      "====> Epoch: 110 Average loss: 2.7469\n",
      "====> Epoch: 120 Average loss: 2.7465\n",
      "====> Epoch: 130 Average loss: 2.7471\n",
      "====> Epoch: 140 Average loss: 2.7476\n",
      "====> Epoch: 150 Average loss: 2.7471\n",
      "====> Epoch: 160 Average loss: 2.7472\n",
      "====> Epoch: 170 Average loss: 2.7458\n",
      "====> Epoch: 180 Average loss: 2.7488\n",
      "====> Epoch: 190 Average loss: 2.7476\n",
      "====> Epoch: 200 Average loss: 2.7464\n",
      "====> Epoch: 210 Average loss: 2.7458\n",
      "====> Epoch: 220 Average loss: 2.7482\n",
      "====> Epoch: 230 Average loss: 2.7456\n",
      "====> Epoch: 240 Average loss: 2.7482\n",
      "====> Epoch: 250 Average loss: 2.7488\n",
      "====> Epoch: 260 Average loss: 2.7452\n",
      "====> Epoch: 270 Average loss: 2.7453\n",
      "====> Epoch: 280 Average loss: 2.7468\n",
      "====> Epoch: 290 Average loss: 2.7439\n",
      "====> Epoch: 300 Average loss: 2.7469\n",
      "====> Epoch: 310 Average loss: 2.7475\n",
      "====> Epoch: 320 Average loss: 2.7458\n",
      "====> Epoch: 330 Average loss: 2.7481\n",
      "====> Epoch: 340 Average loss: 2.7474\n",
      "====> Epoch: 350 Average loss: 2.7470\n",
      "====> Epoch: 360 Average loss: 2.7465\n",
      "====> Epoch: 370 Average loss: 2.7474\n",
      "====> Epoch: 380 Average loss: 2.7462\n",
      "====> Epoch: 390 Average loss: 2.7470\n",
      "====> Epoch: 400 Average loss: 2.7459\n",
      "====> Epoch: 410 Average loss: 2.7473\n",
      "====> Epoch: 420 Average loss: 2.7468\n",
      "====> Epoch: 430 Average loss: 2.7469\n",
      "====> Epoch: 440 Average loss: 2.7476\n",
      "====> Epoch: 450 Average loss: 2.7479\n",
      "====> Epoch: 460 Average loss: 2.7459\n",
      "====> Epoch: 470 Average loss: 2.7477\n",
      "====> Epoch: 480 Average loss: 2.7468\n",
      "====> Epoch: 490 Average loss: 2.7463\n",
      "====> Epoch: 500 Average loss: 2.7478\n",
      "====> Epoch: 510 Average loss: 2.7464\n",
      "====> Epoch: 520 Average loss: 2.7469\n",
      "====> Epoch: 530 Average loss: 2.7461\n",
      "====> Epoch: 540 Average loss: 2.7473\n",
      "====> Epoch: 550 Average loss: 2.7463\n",
      "====> Epoch: 560 Average loss: 2.7481\n",
      "====> Epoch: 570 Average loss: 2.7471\n",
      "====> Epoch: 580 Average loss: 2.7461\n",
      "====> Epoch: 590 Average loss: 2.7477\n",
      "====> Epoch: 600 Average loss: 2.7472\n",
      "====> Epoch: 610 Average loss: 2.7452\n",
      "====> Epoch: 620 Average loss: 2.7480\n",
      "====> Epoch: 630 Average loss: 2.7476\n",
      "====> Epoch: 640 Average loss: 2.7462\n",
      "====> Epoch: 650 Average loss: 2.7465\n",
      "====> Epoch: 660 Average loss: 2.7460\n",
      "====> Epoch: 670 Average loss: 2.7456\n",
      "====> Epoch: 680 Average loss: 2.7451\n",
      "====> Epoch: 690 Average loss: 2.7481\n",
      "====> Epoch: 700 Average loss: 2.7484\n",
      "====> Epoch: 710 Average loss: 2.7467\n",
      "====> Epoch: 720 Average loss: 2.7464\n",
      "====> Epoch: 730 Average loss: 2.7466\n",
      "====> Epoch: 740 Average loss: 2.7475\n",
      "====> Epoch: 750 Average loss: 2.7467\n",
      "====> Epoch: 760 Average loss: 2.7477\n",
      "====> Epoch: 770 Average loss: 2.7453\n",
      "====> Epoch: 780 Average loss: 2.7457\n",
      "====> Epoch: 790 Average loss: 2.7478\n",
      "====> Epoch: 800 Average loss: 2.7481\n",
      "====> Epoch: 810 Average loss: 2.7480\n",
      "====> Epoch: 820 Average loss: 2.7461\n",
      "====> Epoch: 830 Average loss: 2.7461\n",
      "====> Epoch: 840 Average loss: 2.7465\n",
      "====> Epoch: 850 Average loss: 2.7474\n",
      "====> Epoch: 860 Average loss: 2.7476\n",
      "====> Epoch: 870 Average loss: 2.7466\n",
      "====> Epoch: 880 Average loss: 2.7476\n",
      "====> Epoch: 890 Average loss: 2.7488\n",
      "====> Epoch: 900 Average loss: 2.7456\n",
      "====> Epoch: 910 Average loss: 2.7475\n",
      "====> Epoch: 920 Average loss: 2.7466\n",
      "====> Epoch: 930 Average loss: 2.7477\n",
      "====> Epoch: 940 Average loss: 2.7473\n",
      "====> Epoch: 950 Average loss: 2.7466\n",
      "====> Epoch: 960 Average loss: 2.7472\n",
      "====> Epoch: 970 Average loss: 2.7463\n",
      "====> Epoch: 980 Average loss: 2.7472\n",
      "====> Epoch: 990 Average loss: 2.7481\n",
      "====> Epoch: 1000 Average loss: 2.7493\n",
      "====> Epoch: 1010 Average loss: 2.7460\n",
      "====> Epoch: 1020 Average loss: 2.7467\n",
      "====> Epoch: 1030 Average loss: 2.7486\n",
      "====> Epoch: 1040 Average loss: 2.7471\n",
      "====> Epoch: 1050 Average loss: 2.7468\n",
      "====> Epoch: 1060 Average loss: 2.7467\n",
      "====> Epoch: 1070 Average loss: 2.7475\n",
      "====> Epoch: 1080 Average loss: 2.7477\n",
      "====> Epoch: 1090 Average loss: 2.7468\n",
      "====> Epoch: 1100 Average loss: 2.7470\n",
      "====> Epoch: 1110 Average loss: 2.7472\n",
      "====> Epoch: 1120 Average loss: 2.7451\n",
      "====> Epoch: 1130 Average loss: 2.7455\n",
      "====> Epoch: 1140 Average loss: 2.7449\n",
      "====> Epoch: 1150 Average loss: 2.7467\n",
      "====> Epoch: 1160 Average loss: 2.7466\n",
      "====> Epoch: 1170 Average loss: 2.7467\n",
      "====> Epoch: 1180 Average loss: 2.7467\n",
      "====> Epoch: 1190 Average loss: 2.7458\n",
      "====> Epoch: 1200 Average loss: 2.7466\n",
      "====> Epoch: 1210 Average loss: 2.7476\n",
      "====> Epoch: 1220 Average loss: 2.7473\n",
      "====> Epoch: 1230 Average loss: 2.7467\n",
      "====> Epoch: 1240 Average loss: 2.7468\n",
      "====> Epoch: 1250 Average loss: 2.7469\n",
      "====> Epoch: 1260 Average loss: 2.7471\n",
      "====> Epoch: 1270 Average loss: 2.7487\n",
      "====> Epoch: 1280 Average loss: 2.7458\n",
      "====> Epoch: 1290 Average loss: 2.7462\n",
      "====> Epoch: 1300 Average loss: 2.7456\n",
      "====> Epoch: 1310 Average loss: 2.7465\n",
      "====> Epoch: 1320 Average loss: 2.7460\n",
      "====> Epoch: 1330 Average loss: 2.7469\n",
      "====> Epoch: 1340 Average loss: 2.7466\n",
      "====> Epoch: 1350 Average loss: 2.7461\n",
      "====> Epoch: 1360 Average loss: 2.7476\n",
      "====> Epoch: 1370 Average loss: 2.7443\n",
      "====> Epoch: 1380 Average loss: 2.7474\n",
      "====> Epoch: 1390 Average loss: 2.7469\n",
      "====> Epoch: 1400 Average loss: 2.7470\n",
      "====> Epoch: 1410 Average loss: 2.7454\n",
      "====> Epoch: 1420 Average loss: 2.7466\n",
      "====> Epoch: 1430 Average loss: 2.7469\n",
      "====> Epoch: 1440 Average loss: 2.7464\n",
      "====> Epoch: 1450 Average loss: 2.7451\n",
      "====> Epoch: 1460 Average loss: 2.7464\n",
      "====> Epoch: 1470 Average loss: 2.7467\n",
      "====> Epoch: 1480 Average loss: 2.7454\n",
      "====> Epoch: 1490 Average loss: 2.7453\n",
      "====> Epoch: 1500 Average loss: 2.7483\n",
      "====> Epoch: 1510 Average loss: 2.7467\n",
      "====> Epoch: 1520 Average loss: 2.7465\n",
      "====> Epoch: 1530 Average loss: 2.7470\n",
      "====> Epoch: 1540 Average loss: 2.7482\n",
      "====> Epoch: 1550 Average loss: 2.7486\n",
      "====> Epoch: 1560 Average loss: 2.7474\n",
      "====> Epoch: 1570 Average loss: 2.7461\n",
      "====> Epoch: 1580 Average loss: 2.7471\n",
      "====> Epoch: 1590 Average loss: 2.7472\n",
      "====> Epoch: 1600 Average loss: 2.7475\n",
      "====> Epoch: 1610 Average loss: 2.7465\n",
      "====> Epoch: 1620 Average loss: 2.7478\n",
      "====> Epoch: 1630 Average loss: 2.7472\n",
      "====> Epoch: 1640 Average loss: 2.7454\n",
      "====> Epoch: 1650 Average loss: 2.7467\n",
      "====> Epoch: 1660 Average loss: 2.7468\n",
      "====> Epoch: 1670 Average loss: 2.7463\n",
      "====> Epoch: 1680 Average loss: 2.7476\n",
      "====> Epoch: 1690 Average loss: 2.7457\n",
      "====> Epoch: 1700 Average loss: 2.7476\n",
      "====> Epoch: 1710 Average loss: 2.7480\n",
      "====> Epoch: 1720 Average loss: 2.7476\n",
      "====> Epoch: 1730 Average loss: 2.7463\n",
      "====> Epoch: 1740 Average loss: 2.7466\n",
      "====> Epoch: 1750 Average loss: 2.7461\n",
      "====> Epoch: 1760 Average loss: 2.7479\n",
      "====> Epoch: 1770 Average loss: 2.7467\n",
      "====> Epoch: 1780 Average loss: 2.7449\n",
      "====> Epoch: 1790 Average loss: 2.7479\n",
      "====> Epoch: 1800 Average loss: 2.7467\n",
      "====> Epoch: 1810 Average loss: 2.7462\n",
      "====> Epoch: 1820 Average loss: 2.7473\n",
      "====> Epoch: 1830 Average loss: 2.7482\n",
      "====> Epoch: 1840 Average loss: 2.7467\n",
      "====> Epoch: 1850 Average loss: 2.7472\n",
      "====> Epoch: 1860 Average loss: 2.7453\n",
      "====> Epoch: 1870 Average loss: 2.7460\n",
      "====> Epoch: 1880 Average loss: 2.7460\n",
      "====> Epoch: 1890 Average loss: 2.7467\n",
      "====> Epoch: 1900 Average loss: 2.7466\n",
      "====> Epoch: 1910 Average loss: 2.7480\n",
      "====> Epoch: 1920 Average loss: 2.7453\n",
      "====> Epoch: 1930 Average loss: 2.7476\n",
      "====> Epoch: 1940 Average loss: 2.7467\n",
      "====> Epoch: 1950 Average loss: 2.7462\n",
      "====> Epoch: 1960 Average loss: 2.7470\n",
      "====> Epoch: 1970 Average loss: 2.7476\n",
      "====> Epoch: 1980 Average loss: 2.7461\n",
      "====> Epoch: 1990 Average loss: 2.7478\n",
      "====> Epoch: 2000 Average loss: 2.7469\n",
      "====> Epoch: 2010 Average loss: 2.7487\n",
      "====> Epoch: 2020 Average loss: 2.7467\n",
      "====> Epoch: 2030 Average loss: 2.7484\n",
      "====> Epoch: 2040 Average loss: 2.7470\n",
      "====> Epoch: 2050 Average loss: 2.7463\n",
      "====> Epoch: 2060 Average loss: 2.7479\n",
      "====> Epoch: 2070 Average loss: 2.7457\n",
      "====> Epoch: 2080 Average loss: 2.7465\n",
      "====> Epoch: 2090 Average loss: 2.7460\n",
      "====> Epoch: 2100 Average loss: 2.7455\n",
      "====> Epoch: 2110 Average loss: 2.7479\n",
      "====> Epoch: 2120 Average loss: 2.7474\n",
      "====> Epoch: 2130 Average loss: 2.7466\n",
      "====> Epoch: 2140 Average loss: 2.7470\n",
      "====> Epoch: 2150 Average loss: 2.7462\n",
      "====> Epoch: 2160 Average loss: 2.7475\n",
      "====> Epoch: 2170 Average loss: 2.7463\n",
      "====> Epoch: 2180 Average loss: 2.7482\n",
      "====> Epoch: 2190 Average loss: 2.7474\n",
      "====> Epoch: 2200 Average loss: 2.7465\n",
      "====> Epoch: 2210 Average loss: 2.7477\n",
      "====> Epoch: 2220 Average loss: 2.7479\n",
      "====> Epoch: 2230 Average loss: 2.7457\n",
      "====> Epoch: 2240 Average loss: 2.7465\n",
      "====> Epoch: 2250 Average loss: 2.7473\n",
      "====> Epoch: 2260 Average loss: 2.7466\n",
      "====> Epoch: 2270 Average loss: 2.7464\n",
      "====> Epoch: 2280 Average loss: 2.7476\n",
      "====> Epoch: 2290 Average loss: 2.7482\n",
      "====> Epoch: 2300 Average loss: 2.7469\n",
      "====> Epoch: 2310 Average loss: 2.7475\n",
      "====> Epoch: 2320 Average loss: 2.7470\n",
      "====> Epoch: 2330 Average loss: 2.7460\n",
      "====> Epoch: 2340 Average loss: 2.7460\n",
      "====> Epoch: 2350 Average loss: 2.7490\n",
      "====> Epoch: 2360 Average loss: 2.7478\n",
      "====> Epoch: 2370 Average loss: 2.7463\n",
      "====> Epoch: 2380 Average loss: 2.7465\n",
      "====> Epoch: 2390 Average loss: 2.7479\n",
      "====> Epoch: 2400 Average loss: 2.7474\n",
      "====> Epoch: 2410 Average loss: 2.7470\n",
      "====> Epoch: 2420 Average loss: 2.7465\n",
      "====> Epoch: 2430 Average loss: 2.7462\n",
      "====> Epoch: 2440 Average loss: 2.7467\n",
      "====> Epoch: 2450 Average loss: 2.7446\n",
      "====> Epoch: 2460 Average loss: 2.7473\n",
      "====> Epoch: 2470 Average loss: 2.7469\n",
      "====> Epoch: 2480 Average loss: 2.7453\n",
      "====> Epoch: 2490 Average loss: 2.7464\n",
      "====> Epoch: 2500 Average loss: 2.7467\n",
      "====> Epoch: 2510 Average loss: 2.7453\n",
      "====> Epoch: 2520 Average loss: 2.7467\n",
      "====> Epoch: 2530 Average loss: 2.7484\n",
      "====> Epoch: 2540 Average loss: 2.7480\n",
      "====> Epoch: 2550 Average loss: 2.7477\n",
      "====> Epoch: 2560 Average loss: 2.7454\n",
      "====> Epoch: 2570 Average loss: 2.7465\n",
      "====> Epoch: 2580 Average loss: 2.7450\n",
      "====> Epoch: 2590 Average loss: 2.7473\n",
      "====> Epoch: 2600 Average loss: 2.7459\n",
      "====> Epoch: 2610 Average loss: 2.7467\n",
      "====> Epoch: 2620 Average loss: 2.7472\n",
      "====> Epoch: 2630 Average loss: 2.7461\n",
      "====> Epoch: 2640 Average loss: 2.7484\n",
      "====> Epoch: 2650 Average loss: 2.7447\n",
      "====> Epoch: 2660 Average loss: 2.7447\n",
      "====> Epoch: 2670 Average loss: 2.7489\n",
      "====> Epoch: 2680 Average loss: 2.7470\n",
      "====> Epoch: 2690 Average loss: 2.7451\n",
      "====> Epoch: 2700 Average loss: 2.7467\n",
      "====> Epoch: 2710 Average loss: 2.7473\n",
      "====> Epoch: 2720 Average loss: 2.7478\n",
      "====> Epoch: 2730 Average loss: 2.7479\n",
      "====> Epoch: 2740 Average loss: 2.7459\n",
      "====> Epoch: 2750 Average loss: 2.7481\n",
      "====> Epoch: 2760 Average loss: 2.7468\n",
      "====> Epoch: 2770 Average loss: 2.7472\n",
      "====> Epoch: 2780 Average loss: 2.7461\n",
      "====> Epoch: 2790 Average loss: 2.7458\n",
      "====> Epoch: 2800 Average loss: 2.7464\n",
      "====> Epoch: 2810 Average loss: 2.7473\n",
      "====> Epoch: 2820 Average loss: 2.7482\n",
      "====> Epoch: 2830 Average loss: 2.7468\n",
      "====> Epoch: 2840 Average loss: 2.7464\n",
      "====> Epoch: 2850 Average loss: 2.7461\n",
      "====> Epoch: 2860 Average loss: 2.7465\n",
      "====> Epoch: 2870 Average loss: 2.7475\n",
      "====> Epoch: 2880 Average loss: 2.7466\n",
      "====> Epoch: 2890 Average loss: 2.7468\n",
      "====> Epoch: 2900 Average loss: 2.7460\n",
      "====> Epoch: 2910 Average loss: 2.7459\n",
      "====> Epoch: 2920 Average loss: 2.7474\n",
      "====> Epoch: 2930 Average loss: 2.7486\n",
      "====> Epoch: 2940 Average loss: 2.7476\n",
      "====> Epoch: 2950 Average loss: 2.7488\n",
      "====> Epoch: 2960 Average loss: 2.7467\n",
      "====> Epoch: 2970 Average loss: 2.7458\n",
      "====> Epoch: 2980 Average loss: 2.7466\n",
      "====> Epoch: 2990 Average loss: 2.7476\n",
      "====> Epoch: 3000 Average loss: 2.7466\n",
      "====> Epoch: 3010 Average loss: 2.7478\n",
      "====> Epoch: 3020 Average loss: 2.7468\n",
      "====> Epoch: 3030 Average loss: 2.7472\n",
      "====> Epoch: 3040 Average loss: 2.7465\n",
      "====> Epoch: 3050 Average loss: 2.7460\n",
      "====> Epoch: 3060 Average loss: 2.7470\n",
      "====> Epoch: 3070 Average loss: 2.7465\n",
      "====> Epoch: 3080 Average loss: 2.7466\n",
      "====> Epoch: 3090 Average loss: 2.7474\n",
      "====> Epoch: 3100 Average loss: 2.7472\n",
      "====> Epoch: 3110 Average loss: 2.7477\n",
      "====> Epoch: 3120 Average loss: 2.7463\n",
      "====> Epoch: 3130 Average loss: 2.7475\n",
      "====> Epoch: 3140 Average loss: 2.7459\n",
      "====> Epoch: 3150 Average loss: 2.7455\n",
      "====> Epoch: 3160 Average loss: 2.7455\n",
      "====> Epoch: 3170 Average loss: 2.7467\n",
      "====> Epoch: 3180 Average loss: 2.7485\n",
      "====> Epoch: 3190 Average loss: 2.7464\n",
      "====> Epoch: 3200 Average loss: 2.7477\n",
      "====> Epoch: 3210 Average loss: 2.7468\n",
      "====> Epoch: 3220 Average loss: 2.7478\n",
      "====> Epoch: 3230 Average loss: 2.7462\n",
      "====> Epoch: 3240 Average loss: 2.7467\n",
      "====> Epoch: 3250 Average loss: 2.7461\n",
      "====> Epoch: 3260 Average loss: 2.7473\n",
      "====> Epoch: 3270 Average loss: 2.7466\n",
      "====> Epoch: 3280 Average loss: 2.7472\n",
      "====> Epoch: 3290 Average loss: 2.7466\n",
      "====> Epoch: 3300 Average loss: 2.7472\n",
      "====> Epoch: 3310 Average loss: 2.7458\n",
      "====> Epoch: 3320 Average loss: 2.7447\n",
      "====> Epoch: 3330 Average loss: 2.7460\n",
      "====> Epoch: 3340 Average loss: 2.7461\n",
      "====> Epoch: 3350 Average loss: 2.7459\n",
      "====> Epoch: 3360 Average loss: 2.7467\n",
      "====> Epoch: 3370 Average loss: 2.7471\n",
      "====> Epoch: 3380 Average loss: 2.7468\n",
      "====> Epoch: 3390 Average loss: 2.7461\n",
      "====> Epoch: 3400 Average loss: 2.7464\n",
      "====> Epoch: 3410 Average loss: 2.7459\n",
      "====> Epoch: 3420 Average loss: 2.7465\n",
      "====> Epoch: 3430 Average loss: 2.7459\n",
      "====> Epoch: 3440 Average loss: 2.7477\n",
      "====> Epoch: 3450 Average loss: 2.7474\n",
      "====> Epoch: 3460 Average loss: 2.7475\n",
      "====> Epoch: 3470 Average loss: 2.7464\n",
      "====> Epoch: 3480 Average loss: 2.7481\n",
      "====> Epoch: 3490 Average loss: 2.7486\n",
      "====> Epoch: 3500 Average loss: 2.7461\n",
      "====> Epoch: 3510 Average loss: 2.7472\n",
      "====> Epoch: 3520 Average loss: 2.7483\n",
      "====> Epoch: 3530 Average loss: 2.7472\n",
      "====> Epoch: 3540 Average loss: 2.7452\n",
      "====> Epoch: 3550 Average loss: 2.7462\n",
      "====> Epoch: 3560 Average loss: 2.7473\n",
      "====> Epoch: 3570 Average loss: 2.7470\n",
      "====> Epoch: 3580 Average loss: 2.7471\n",
      "====> Epoch: 3590 Average loss: 2.7484\n",
      "====> Epoch: 3600 Average loss: 2.7470\n",
      "====> Epoch: 3610 Average loss: 2.7464\n",
      "====> Epoch: 3620 Average loss: 2.7474\n",
      "====> Epoch: 3630 Average loss: 2.7462\n",
      "====> Epoch: 3640 Average loss: 2.7475\n",
      "====> Epoch: 3650 Average loss: 2.7461\n",
      "====> Epoch: 3660 Average loss: 2.7478\n",
      "====> Epoch: 3670 Average loss: 2.7467\n",
      "====> Epoch: 3680 Average loss: 2.7464\n",
      "====> Epoch: 3690 Average loss: 2.7465\n",
      "====> Epoch: 3700 Average loss: 2.7473\n",
      "====> Epoch: 3710 Average loss: 2.7485\n",
      "====> Epoch: 3720 Average loss: 2.7473\n",
      "====> Epoch: 3730 Average loss: 2.7474\n",
      "====> Epoch: 3740 Average loss: 2.7459\n",
      "====> Epoch: 3750 Average loss: 2.7471\n",
      "====> Epoch: 3760 Average loss: 2.7458\n",
      "====> Epoch: 3770 Average loss: 2.7464\n",
      "====> Epoch: 3780 Average loss: 2.7476\n",
      "====> Epoch: 3790 Average loss: 2.7481\n",
      "====> Epoch: 3800 Average loss: 2.7457\n",
      "====> Epoch: 3810 Average loss: 2.7460\n",
      "====> Epoch: 3820 Average loss: 2.7478\n",
      "====> Epoch: 3830 Average loss: 2.7474\n",
      "====> Epoch: 3840 Average loss: 2.7462\n",
      "====> Epoch: 3850 Average loss: 2.7462\n",
      "====> Epoch: 3860 Average loss: 2.7469\n",
      "====> Epoch: 3870 Average loss: 2.7449\n",
      "====> Epoch: 3880 Average loss: 2.7481\n",
      "====> Epoch: 3890 Average loss: 2.7464\n",
      "====> Epoch: 3900 Average loss: 2.7458\n",
      "====> Epoch: 3910 Average loss: 2.7479\n",
      "====> Epoch: 3920 Average loss: 2.7483\n",
      "====> Epoch: 3930 Average loss: 2.7459\n",
      "====> Epoch: 3940 Average loss: 2.7467\n",
      "====> Epoch: 3950 Average loss: 2.7488\n",
      "====> Epoch: 3960 Average loss: 2.7460\n",
      "====> Epoch: 3970 Average loss: 2.7464\n",
      "====> Epoch: 3980 Average loss: 2.7468\n",
      "====> Epoch: 3990 Average loss: 2.7476\n",
      "====> Epoch: 4000 Average loss: 2.7458\n",
      "====> Epoch: 4010 Average loss: 2.7474\n",
      "====> Epoch: 4020 Average loss: 2.7468\n",
      "====> Epoch: 4030 Average loss: 2.7482\n",
      "====> Epoch: 4040 Average loss: 2.7459\n",
      "====> Epoch: 4050 Average loss: 2.7493\n",
      "====> Epoch: 4060 Average loss: 2.7465\n",
      "====> Epoch: 4070 Average loss: 2.7461\n",
      "====> Epoch: 4080 Average loss: 2.7475\n",
      "====> Epoch: 4090 Average loss: 2.7471\n",
      "====> Epoch: 4100 Average loss: 2.7462\n",
      "====> Epoch: 4110 Average loss: 2.7461\n",
      "====> Epoch: 4120 Average loss: 2.7461\n",
      "====> Epoch: 4130 Average loss: 2.7463\n",
      "====> Epoch: 4140 Average loss: 2.7467\n",
      "====> Epoch: 4150 Average loss: 2.7455\n",
      "====> Epoch: 4160 Average loss: 2.7465\n",
      "====> Epoch: 4170 Average loss: 2.7465\n",
      "====> Epoch: 4180 Average loss: 2.7469\n",
      "====> Epoch: 4190 Average loss: 2.7473\n",
      "====> Epoch: 4200 Average loss: 2.7474\n",
      "====> Epoch: 4210 Average loss: 2.7462\n",
      "====> Epoch: 4220 Average loss: 2.7461\n",
      "====> Epoch: 4230 Average loss: 2.7481\n",
      "====> Epoch: 4240 Average loss: 2.7465\n",
      "====> Epoch: 4250 Average loss: 2.7478\n",
      "====> Epoch: 4260 Average loss: 2.7468\n",
      "====> Epoch: 4270 Average loss: 2.7465\n",
      "====> Epoch: 4280 Average loss: 2.7493\n",
      "====> Epoch: 4290 Average loss: 2.7465\n",
      "====> Epoch: 4300 Average loss: 2.7469\n",
      "====> Epoch: 4310 Average loss: 2.7466\n",
      "====> Epoch: 4320 Average loss: 2.7454\n",
      "====> Epoch: 4330 Average loss: 2.7450\n",
      "====> Epoch: 4340 Average loss: 2.7469\n",
      "====> Epoch: 4350 Average loss: 2.7459\n",
      "====> Epoch: 4360 Average loss: 2.7479\n",
      "====> Epoch: 4370 Average loss: 2.7462\n",
      "====> Epoch: 4380 Average loss: 2.7487\n",
      "====> Epoch: 4390 Average loss: 2.7470\n",
      "====> Epoch: 4400 Average loss: 2.7468\n",
      "====> Epoch: 4410 Average loss: 2.7461\n",
      "====> Epoch: 4420 Average loss: 2.7462\n",
      "====> Epoch: 4430 Average loss: 2.7475\n",
      "====> Epoch: 4440 Average loss: 2.7474\n",
      "====> Epoch: 4450 Average loss: 2.7476\n",
      "====> Epoch: 4460 Average loss: 2.7479\n",
      "====> Epoch: 4470 Average loss: 2.7465\n",
      "====> Epoch: 4480 Average loss: 2.7454\n",
      "====> Epoch: 4490 Average loss: 2.7479\n",
      "====> Epoch: 4500 Average loss: 2.7457\n",
      "====> Epoch: 4510 Average loss: 2.7460\n",
      "====> Epoch: 4520 Average loss: 2.7468\n",
      "====> Epoch: 4530 Average loss: 2.7449\n",
      "====> Epoch: 4540 Average loss: 2.7466\n",
      "====> Epoch: 4550 Average loss: 2.7470\n",
      "====> Epoch: 4560 Average loss: 2.7477\n",
      "====> Epoch: 4570 Average loss: 2.7482\n",
      "====> Epoch: 4580 Average loss: 2.7472\n",
      "====> Epoch: 4590 Average loss: 2.7479\n",
      "====> Epoch: 4600 Average loss: 2.7457\n",
      "====> Epoch: 4610 Average loss: 2.7470\n",
      "====> Epoch: 4620 Average loss: 2.7459\n",
      "====> Epoch: 4630 Average loss: 2.7465\n",
      "====> Epoch: 4640 Average loss: 2.7476\n",
      "====> Epoch: 4650 Average loss: 2.7486\n",
      "====> Epoch: 4660 Average loss: 2.7486\n",
      "====> Epoch: 4670 Average loss: 2.7477\n",
      "====> Epoch: 4680 Average loss: 2.7474\n",
      "====> Epoch: 4690 Average loss: 2.7486\n",
      "====> Epoch: 4700 Average loss: 2.7469\n",
      "====> Epoch: 4710 Average loss: 2.7472\n",
      "====> Epoch: 4720 Average loss: 2.7468\n",
      "====> Epoch: 4730 Average loss: 2.7455\n",
      "====> Epoch: 4740 Average loss: 2.7473\n",
      "====> Epoch: 4750 Average loss: 2.7476\n",
      "====> Epoch: 4760 Average loss: 2.7468\n",
      "====> Epoch: 4770 Average loss: 2.7464\n",
      "====> Epoch: 4780 Average loss: 2.7477\n",
      "====> Epoch: 4790 Average loss: 2.7470\n",
      "====> Epoch: 4800 Average loss: 2.7462\n",
      "====> Epoch: 4810 Average loss: 2.7468\n",
      "====> Epoch: 4820 Average loss: 2.7464\n",
      "====> Epoch: 4830 Average loss: 2.7479\n",
      "====> Epoch: 4840 Average loss: 2.7471\n",
      "====> Epoch: 4850 Average loss: 2.7460\n",
      "====> Epoch: 4860 Average loss: 2.7481\n",
      "====> Epoch: 4870 Average loss: 2.7464\n",
      "====> Epoch: 4880 Average loss: 2.7467\n",
      "====> Epoch: 4890 Average loss: 2.7468\n",
      "====> Epoch: 4900 Average loss: 2.7466\n",
      "====> Epoch: 4910 Average loss: 2.7462\n",
      "====> Epoch: 4920 Average loss: 2.7489\n",
      "====> Epoch: 4930 Average loss: 2.7456\n",
      "====> Epoch: 4940 Average loss: 2.7479\n",
      "====> Epoch: 4950 Average loss: 2.7480\n",
      "====> Epoch: 4960 Average loss: 2.7456\n",
      "====> Epoch: 4970 Average loss: 2.7480\n",
      "====> Epoch: 4980 Average loss: 2.7481\n",
      "====> Epoch: 4990 Average loss: 2.7480\n",
      "====> Epoch: 5000 Average loss: 2.7471\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 18092 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17424 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17220 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 14860 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 15016 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 11276 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16644 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "optimizer = optim.Adam(vae.parameters())\n",
    "# return reconstruction error + KL divergence losses\n",
    "def loss_function(recon_x, x, mu, log_var):\n",
    "    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')\n",
    "    KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())\n",
    "    return BCE + KLD\n",
    "\n",
    "sizes = [300] \n",
    "sampling_noise = 0.05 \n",
    "loss_across_sizes = []\n",
    "for i in range(len(sizes)):\n",
    "    totalnum = sizes[i]\n",
    "    indices = torch.arange(totalnum)\n",
    "    tr_10k = data_utils.Subset(ds, indices)\n",
    "    dl_10k = torch.utils.data.DataLoader(dataset=tr_10k, batch_size=bs, shuffle=True)\n",
    "    vae = VAE()\n",
    "    vae.to(device)\n",
    "\n",
    "    def train(epoch):\n",
    "        vae.train()\n",
    "        train_loss = 0\n",
    "        for batch_idx, (data, _) in enumerate(dl_10k): # was train_loader\n",
    "            data = data.cuda()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            recon_batch, mu, log_var = vae(data)\n",
    "            loss = loss_function(recon_batch, data, mu, log_var)\n",
    "\n",
    "            loss.backward()\n",
    "            train_loss += loss.item()\n",
    "            optimizer.step()\n",
    "        if epoch % 10 == 0:\n",
    "            print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, train_loss / len(train_loader.dataset)))\n",
    "    def test():\n",
    "        vae.eval()\n",
    "        test_loss= 0\n",
    "        with torch.no_grad():\n",
    "            for data, _ in test_loader:\n",
    "                data = data.cuda()\n",
    "                recon, mu, log_var = vae(data)\n",
    "\n",
    "                # sum up batch loss\n",
    "                test_loss += loss_function(recon, data, mu, log_var).item()\n",
    "\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        print('====> Test set loss: {:.4f}'.format(test_loss))\n",
    "    for epoch in range(1, 5001): #51):\n",
    "        train(epoch)\n",
    "        #test()\n",
    "        \n",
    "    # Evaluate via denoising loss!\n",
    "    test_sample_func_train, test_sample_func_test = make_test_sample_functions(dl_iterator=iter(dl), dl_test_iterator=iter(dl_test))\n",
    "    m = 784\n",
    "    k = 20\n",
    "    def add_noise(invec):\n",
    "        with torch.no_grad():\n",
    "            return invec + sampling_noise*torch.randn((P, m), device=device)\n",
    "\n",
    "    sampler = compose2(add_noise, test_sample_func_train)\n",
    "    add_noise_func = add_noise\n",
    "    test_sample_func = test_sample_func_test\n",
    "\n",
    "    iters = 1000\n",
    "    test_losses = []\n",
    "    num_batch = 5\n",
    "    lr = 1e-2\n",
    "\n",
    "    test_losses, ys_noiseless, outputs, ys = lightweight_denoising_loss(num_batch, lr, iters, P, k, vae.decode, test_sample_func, add_noise_func, do_plots=True)\n",
    "    loss_across_sizes.append(torch.mean(torch.tensor(test_losses)))\n",
    "    with torch.no_grad():\n",
    "        z = torch.randn(bs, 20).cuda()\n",
    "        sample = vae.decode(z).cuda()\n",
    "        plt.imshow(sample[0,...].view(28,28).cpu())\n",
    "        plt.title('num datapoints' + str(totalnum))\n",
    "        plt.show()\n",
    "        \n",
    "        fig, axs = plt.subplots(1, 3)\n",
    "        axs[0].imshow(ys_noiseless[0,...].reshape(28,28).cpu())\n",
    "        axs[0].set_title('noiseless')\n",
    "        axs[1].imshow(outputs[0,...].reshape(28,28).detach().cpu())\n",
    "        axs[1].set_title('recon')\n",
    "        axs[2].imshow(ys[0,...].reshape(28,28).cpu())\n",
    "        axs[2].set_title('noisy')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11932 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ys.shape\n",
    "vae.decode(vae.encode(ys)[0]).shape\n",
    "outs = vae.decode(vae.encode(ys)[0])\n",
    "plt.imshow(outs[0,...].detach().reshape(28,28).cpu())\n",
    "plt.title('Using encoder for inversion')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cryo",
   "language": "python",
   "name": "cryo"
  },
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
