03_03_vae_digits_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VariationalAutoencoder(\n",
       "  (enc_conv_layers): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (4): Flatten()\n",
       "    (5): TwoTupleLinear(\n",
       "      (lin1): Linear(in_features=2304, out_features=2, bias=True)\n",
       "      (lin2): Linear(in_features=2304, out_features=2, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (lamdba_sampling_layer): Lambda()\n",
       "  (dec_conv_layers): Sequential(\n",
       "    (0): Linear(in_features=2, out_features=2304, bias=True)\n",
       "    (1): View()\n",
       "    (2): Sequential(\n",
       "      (0): ConvTranspose2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): ConvTranspose2d(64, 64, kernel_size=(3, 3), stride=(2, 2), output_padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (4): Sequential(\n",
       "      (0): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "      (2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Dropout(p=0.25)\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): ConvTranspose2d(32, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): Sigmoid()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch as t\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm\n",
    "from models.VariationalAutoencoder import VariationalAutoencoder\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "bs = 128\n",
    "train_ds = datasets.MNIST(root='./data/', train=True, transform=transforms.ToTensor(), download=True)\n",
    "train_dl = t.utils.data.DataLoader(dataset=train_ds, batch_size=bs, shuffle=True, drop_last=True)\n",
    "\n",
    "device = t.device('cuda') if t.cuda.is_available() else 'cpu'\n",
    "model = VariationalAutoencoder(train_ds[0][0][None], in_c=1, enc_out_c=[32, 64, 64, 64],\n",
    "                    enc_ks=[3, 3, 3, 3], enc_pads=[1, 1, 0, 1], enc_strides=[1, 2, 2, 1],\n",
    "                    dec_out_c=[64, 64, 32, 1], dec_ks=[3, 3, 3, 3], dec_strides=[1, 2, 2, 1],\n",
    "                    dec_pads=[1, 0, 1, 1], dec_op_pads=[0, 1, 1, 0], z_dim=2)\n",
    "model.cuda(device)\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vae_kl_loss(mu, log_var):\n",
    "    return -.5 * t.sum(1 + log_var - mu ** 2 - log_var.exp())\n",
    "\n",
    "def vae_loss(y_pred, mu, log_var, y_true):\n",
    "    r_loss = F.binary_cross_entropy(y_pred, y_true, reduction='sum')\n",
    "    kl_loss = vae_kl_loss(mu, log_var)\n",
    "    return r_loss + kl_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/20 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(108783.1250, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(62251.4727, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(45925.7930, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(33259.5938, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(29036.8750, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(28326.6035, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(25913.1426, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(25414.1348, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(25909.0078, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(25462.0059, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(24686.7383, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(23606.9512, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22902.1191, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(24357.7031, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(23244.5645, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 1/20 [00:19<06:19, 19.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(23037.9980, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(23119.1738, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(24233.2500, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22439.0488, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(23018.7676, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22614.6309, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21595.1387, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22400.3887, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22202.3477, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22966.3867, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22581.9355, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22255.6836, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21529.6387, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22259.2793, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21393.3711, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 2/20 [00:40<06:01, 20.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(22190.8906, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22123.2051, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21594.9023, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22594.9883, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22748.5391, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22140.2344, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21788.3027, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22641.7832, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22495.2266, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22827.6992, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21153.1855, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22124.2305, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21570.4609, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21029.9238, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22264.3359, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▌        | 3/20 [01:00<05:42, 20.13s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21611.2695, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22143.6465, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22589.4160, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22484.4570, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22246.1875, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22359.3848, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21270.3359, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22727.1895, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21054.2695, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22352.3887, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21217.8848, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21564.9297, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21843.3418, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22552.9492, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21861.8965, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 4/20 [01:20<05:23, 20.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20933.6582, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21949.1719, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21991.1328, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20884.3711, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21483.3809, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22086.7715, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21267.1152, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21622.9355, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20439.7559, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22263.3770, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21221.8828, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21802.8086, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21720.2578, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22179.9336, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20663.1641, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 5/20 [01:41<05:03, 20.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21559.5840, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22061.6699, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21153.0215, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21486.0391, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21442.1445, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21496.2734, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21192.7500, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20276.7148, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21352.6270, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22349.9102, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21350.8496, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22205.0820, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21442.8203, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21502.7949, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21091.8223, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 6/20 [02:01<04:43, 20.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20751.1191, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20776.8750, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21333.7715, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22407.3594, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20934.7461, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21431.1348, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21933.2188, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21241.0781, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21304.9160, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21859.6426, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22275.0859, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21586.7383, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21066.2090, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21571.2168, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▌      | 7/20 [02:21<04:23, 20.27s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21386.5508, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21293.1387, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20630.8770, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21602.8613, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20748.3301, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20933.3730, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21120.1914, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20769.3887, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20791.7754, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21780.4297, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21240.8789, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21318.3203, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21007.9141, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20785.2480, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21631.5234, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20859.2012, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 8/20 [02:41<04:00, 20.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20848.4727, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21647.4199, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21677.4375, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21621.7930, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21519.9277, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19765.1934, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21992.1543, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20386.3633, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21990.9531, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21588.9941, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21022.2578, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20993.9746, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21107.9883, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20819.5078, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▌     | 9/20 [02:59<03:35, 19.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21611.8086, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20840.5566, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21583.4512, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21567.5879, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21033.4180, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20947.8379, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21062.5605, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20808.1484, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21447.1230, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21250.5332, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21374.8047, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20660.9609, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20597.6289, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21391.2598, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20243.7012, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21114.8105, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 10/20 [03:20<03:18, 19.83s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20337.6133, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21172.9805, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20919.3398, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20578.4844, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21682.1914, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22401.3867, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21647.3535, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20666.1465, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21816.6016, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21849.7891, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21417.2363, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21211.2324, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21360.4727, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20697.5195, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21442.7949, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▌    | 11/20 [03:40<03:00, 20.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(19683.2891, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21137.7871, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21713.8320, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20426.9043, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21045.4395, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21926.2910, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21071.3340, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21130.3691, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20735.4199, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21153.0176, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21998.1758, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21597.0488, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19986.7441, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21480.1055, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21690.0410, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 12/20 [04:01<02:41, 20.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21190.2793, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21245.7402, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21434.6914, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20338.1191, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20583.3008, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21475.6172, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20887.4785, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20442.9941, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20589.5234, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20956.0762, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20453.7090, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20691.4375, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(22521.5547, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20925.5566, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 65%|██████▌   | 13/20 [04:20<02:19, 19.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20641.8926, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21997.5254, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20682.2812, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21205.6250, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20176.5078, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20965.5820, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21621.4863, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21868.8652, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21395.9219, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21787.8145, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21045.1289, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21258.1094, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21427.2734, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20059.7852, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20987.4551, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21017.8145, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 14/20 [04:40<01:59, 19.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20002.5938, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20303.3066, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20412.3672, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20402.8926, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21089.7656, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21429.5586, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20140.8262, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20715.5977, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20989.5312, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21019.5391, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20238.3262, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21267.3828, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21447.5176, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20335.7930, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20634.4531, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 15/20 [05:01<01:40, 20.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21133.9277, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20937.1289, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20694.6074, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21628.4551, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21516.1250, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20644.4160, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21985.0312, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20451.7227, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20565.2773, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20528.4844, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21309.9062, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20333.3945, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20885.6348, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21781.8535, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21145.2969, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 16/20 [05:20<01:20, 20.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21076.9551, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21921.6035, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20530.2617, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20897.3809, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20990.8066, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19573.1523, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20789.1074, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21907.3730, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19669.4199, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21225.5430, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20797.1562, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20979.4609, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20903.4805, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20800.6797, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21609.2773, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▌ | 17/20 [05:41<01:00, 20.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(21177.4102, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20895.1641, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21560.3203, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21400.1094, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19796.9922, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20789.3281, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21294.1621, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20533.0488, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21162.4980, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20856.9160, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20804.6133, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(19932.3945, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20674.0918, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20925.8887, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20478.5273, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 18/20 [06:02<00:40, 20.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20421.5742, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21626.6543, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20620.6426, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21470.6055, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20645.3281, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20720.9844, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21151.9727, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21443.3496, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21731.2148, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20645.4023, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21662.0234, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20963.5645, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20453.7891, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20917.6094, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20432.3555, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 95%|█████████▌| 19/20 [06:22<00:20, 20.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20687.7441, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21570.8555, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20646.7324, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20537.5586, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20609.2578, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20168.6504, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20183.7617, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21104.4941, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21474.7969, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21322.9199, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21733.2480, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20674.8691, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(21726.9375, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20735.9023, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "tensor(20616.7090, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [06:42<00:00, 20.13s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(20704.5410, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "lr = .0005\n",
    "for epoch in tqdm(range(20)):\n",
    "    optimizer = t.optim.Adam(model.parameters(), lr=lr / (epoch * 2 + 1), betas=(.9, .99), weight_decay=1e-2)\n",
    "    for i, (data, _) in enumerate(train_dl):\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        pred, mu, log_var = model(data)\n",
    "        loss = vae_loss(pred, mu, log_var, data)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if i % 33 == 0:\n",
    "            print(loss)\n",
    "\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.save(model.state_dict(), 'models/state_dicts/03_03.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}