03_05_vae_faces_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE Training - Faces dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CelebVariationalAutoencoder(\n",
       "  (enc_conv_layers): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 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, inplace=False)\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, inplace=False)\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, inplace=False)\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, inplace=False)\n",
       "    )\n",
       "    (4): Flatten()\n",
       "    (5): TwoTupleLinear(\n",
       "      (lin1): Linear(in_features=61504, out_features=200, bias=True)\n",
       "      (lin2): Linear(in_features=61504, out_features=200, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (lamdba_sampling_layer): Lambda()\n",
       "  (dec_conv_layers): Sequential(\n",
       "    (0): Linear(in_features=200, out_features=61504, 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, inplace=False)\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, inplace=False)\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, inplace=False)\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): ConvTranspose2d(32, 3, 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",
    "import zipfile\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from models.CelebVariationalAutoencoder import CelebVariationalAutoencoder\n",
    "from utils.data.FaceDataset import FaceDataset\n",
    "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
    "\n",
    "# this will take some time :), its 1.3gb to download\n",
    "faces_zip = 'data/faces.zip'\n",
    "if not os.path.exists(faces_zip):\n",
    "    gdd.download_file_from_google_drive(file_id='0B7EVK8r0v71pZjFTYXZWM3FlRnM',\n",
    "                                        dest_path=faces_zip)\n",
    "    with zipfile.ZipFile(faces_zip, 'r') as zip_ref:\n",
    "        zip_ref.extractall('data/faces')\n",
    "\n",
    "bs = 64\n",
    "train_ds = FaceDataset(\"data/faces/img_align_celeba/\")\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 = CelebVariationalAutoencoder(train_ds[0][0][None], in_c=3, 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, 3], 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=200)\n",
    "model.cuda(device)\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/3 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(55878044., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(42448192., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(35465824., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(30940976., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(30071830., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(29281372., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(29052596., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27535826., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27238960., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27496606., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26781960., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26130110., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25813870., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27273876., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26892190., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26024736., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26949876., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26472106., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26260496., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26153928., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26287676., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26513824., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25956710., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26058470., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25970020., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25535236., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26973130., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26057118., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26779578., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26420958., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25965270., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26218594., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25735708., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26116530., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26904346., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25901428., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26102814., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25515014., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26304100., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26361832., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26705472., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26061162., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26869254., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26898662., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26288678., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25566956., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26319892., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26385262., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26690902., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26284820., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26611424., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26357958., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26843294., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26620976., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25771912., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26885734., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25999406., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26548454., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26381174., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25897450., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25594598., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26320334., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25977850., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26198476., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26814146., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26355902., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26377122., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26323280., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26271604., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25656444., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27384322., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26178608., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25970586., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26607602., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26766154., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26518406., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26297624., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26983868., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25325210., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26309626., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26944704., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27432664., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26377644., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27150242., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26723958., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26018260., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26384850., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27105222., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26216120., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25574092., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26294984., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26635760., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26545340., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26011878., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27147380., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26688866., device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|███▎      | 1/3 [26:47<53:34, 1607.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(26885404., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25962064., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25586096., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26627036., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25364480., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26790918., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26882406., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27293300., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25610908., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26584938., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26097558., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25284448., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26113560., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25607128., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25341408., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26445146., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25692274., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26569818., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25983152., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26016606., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25896182., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26332614., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26532280., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26805112., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27009986., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26069036., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26883172., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26394892., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26656354., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26262866., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26300560., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25920076., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26917910., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26243426., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25981500., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26514712., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26401446., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26582554., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25654798., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26431698., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26210082., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25512402., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26020502., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25638336., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25386702., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25604708., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26255514., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26180970., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25764654., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26254358., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25945592., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26717814., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26437294., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26286122., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25732712., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26482400., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26433004., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26179166., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26672496., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26937354., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26086700., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25831790., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26729776., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26579550., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26656944., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25937376., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26829996., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26283622., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26309688., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25913662., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26730634., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26051656., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26005032., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26534860., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26938402., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26368492., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25867966., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26415420., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26439648., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25837110., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26205986., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26411468., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26475990., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26458860., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25871552., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26319676., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25684070., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26760426., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26056720., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26137878., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26359130., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26061296., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26788874., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25947798., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25281452., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25657542., device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 67%|██████▋   | 2/3 [53:06<26:38, 1598.89s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(26537340., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26078576., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25644586., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25917624., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26084042., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26706180., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25959336., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25217452., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26832634., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26372622., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26084352., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25709944., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26116460., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26044390., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25651988., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26190934., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26602564., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25967130., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25711444., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25630196., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26032288., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26885158., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26252074., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26354918., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25953786., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26182534., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26999398., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25533260., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26220494., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26000078., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25894400., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26036510., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26345436., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25448516., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26297948., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26359906., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25650164., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25615754., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26212704., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26233390., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26640668., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25936128., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26092002., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26242402., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27189866., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26257972., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25960308., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26608982., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26925778., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26038356., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26062814., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25884278., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26309318., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25958544., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26325034., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26200580., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26842620., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25888954., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26540086., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26948732., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26051482., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26544854., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25866096., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25798624., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25854776., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25760192., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25661748., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27045796., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25662870., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26419158., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26487766., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26097896., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25904014., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26426006., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25980722., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26155830., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25863640., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26394876., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25756784., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27241136., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25935446., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26936772., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25812848., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26416574., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27196112., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26154818., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26406502., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25891590., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26222378., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25820694., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26800756., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26672304., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26625590., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26565336., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26885380., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26404840., device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [1:19:28<00:00, 1593.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1657669632., device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "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, r_loss_factor=1000):\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_factor * r_loss + kl_loss\n",
    "\n",
    "lr = .0005\n",
    "for epoch in tqdm(range(3)):\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",
    "        t.nn.utils.clip_grad_norm_(model.parameters(), .25)\n",
    "        optimizer.step()\n",
    "        if i % 33 == 0:\n",
    "            print(loss / bs)\n",
    "\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.save(model.state_dict(), 'models/state_dicts/03_05.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(26122478., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26074100., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25833470., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26289132., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26832608., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26591926., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26524614., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26615118., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27162468., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26599992., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26118604., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26165174., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26279220., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26323964., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26328696., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26620050., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26036438., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25246744., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26836970., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25494254., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26225494., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26293198., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26322108., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26788000., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26184946., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26335274., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26283558., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27065532., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26230236., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26593868., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25674768., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26314638., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26858398., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25940904., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26429268., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25854760., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27097696., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26558422., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27114802., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26557290., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25551676., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26410144., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26448782., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26361246., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26217178., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26430666., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26159936., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26373250., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25818872., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26003754., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25880670., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26245548., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26031422., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25932160., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26389632., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25945974., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26067560., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26077414., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26181338., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26628694., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25921818., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26358414., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26629940., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26760246., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26977702., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26430738., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26241870., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26061948., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26039070., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25459510., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26633108., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25790268., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26630734., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26372720., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25954780., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25631672., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25623856., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26408578., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25937178., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26577644., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26400248., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25848334., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25938550., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26171386., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25679722., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26444286., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26078518., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26240480., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26412424., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26939664., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25792186., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26330858., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26382954., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26302356., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26264756., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26708132., device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 1/2 [26:45<26:45, 1605.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(26394400., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26236288., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25804286., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26903614., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26662452., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26644092., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26373228., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26691236., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25686442., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26356246., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26791182., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25466834., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26669580., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26525412., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26029278., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25963376., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26508208., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26223580., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25858830., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26199308., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26541052., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26433484., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26176724., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26017448., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26222236., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26074008., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26490210., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26064384., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25480158., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26219554., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26386564., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26585404., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26309422., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26849358., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26542350., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25915976., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26135018., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26251626., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25932488., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26904350., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26461708., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25530004., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26248746., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27078402., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25570268., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26136538., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26906690., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25590184., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25954088., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26775660., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26136854., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25659486., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25946472., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26944058., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26299664., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26612890., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26729534., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26856292., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25852402., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26520628., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26312244., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26743934., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26938210., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26499062., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26549576., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26304136., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25425808., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25931600., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25524336., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25462684., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26694384., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27117980., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25768554., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25871436., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26660868., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25756176., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25665548., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(27297642., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26743802., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26446776., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26040122., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26209390., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26227144., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25951220., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26430186., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25947370., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26302596., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26672846., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26035884., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26825328., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25880072., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26364068., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26857122., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26762574., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(26888380., device='cuda:0', grad_fn=<DivBackward0>)\n",
      "tensor(25472702., device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [53:30<00:00, 1605.39s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1662345472., device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "lr = .0005\n",
    "for epoch in tqdm(range(3, 5)):\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",
    "        t.nn.utils.clip_grad_norm_(model.parameters(), .25)\n",
    "        optimizer.step()\n",
    "        if i % 33 == 0:\n",
    "            print(loss / bs)\n",
    "\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.save(model.state_dict(), 'models/state_dict/03_05_full.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(25974148., device='cuda:0', grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss / bs"
   ]
  }
 ],
 "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.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}