Classical Autoencoder-CelebA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "from torchvision.datasets import MNIST\n",
    "import os\n",
    "\n",
    "from functionalities import filemanager as fm\n",
    "from functionalities import dataloader as dl\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from functionalities import loss as lo\n",
    "from functionalities import plot as pl\n",
    "from functionalities import trainer as tr\n",
    "from functionalities import gpu \n",
    "\n",
    "from architecture import CelebA_autoencoder as celeba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pretraining Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epoch = 2 #10\n",
    "batch_size = 32\n",
    "lr_init = 1e-3\n",
    "image_size = 178\n",
    "milestones = [8, 9, 10]\n",
    "latent_dim_lst = [128, 1000, 4000] #[2 ** x for x in range(0, 11, 2)]\n",
    "number_dev = 0\n",
    "get_model = celeba.celeba_autoencoder\n",
    "modelname = \"celeba_classic_bottleneck\"\n",
    "\n",
    "number_dev = 0\n",
    "device = gpu.get_device(number_dev)\n",
    "print(device)\n",
    "\n",
    "\n",
    "# you need to manually download celebA dataset, have the celebA images save in folder ./img_align_celeba/dummy_class/...\n",
    "IMAGE_PATH = './img_align_celeba/'\n",
    "transform = transforms.Compose([\n",
    "    #transforms.Scale(image_size),\n",
    "    #transforms.Resize(image_size),\n",
    "    transforms.ToTensor(),\n",
    "    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "dataset = ImageFolder(IMAGE_PATH, transform)\n",
    "trainloader, testloader = dl.split_dataset(dataset, 0.2, batch_size, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr.train_bottleneck_classic(num_epoch, get_model, modelname, milestones, latent_dim_lst, trainloader,\n",
    "                     None, testloader, lr_init=lr_init, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot Reconstruction and Difference Images Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_img = 25\n",
    "grid_row_size = 5\n",
    "\n",
    "img, label = next(iter(testloader))\n",
    "#img = img.view(img.size(0), -1)\n",
    "img = Variable(img).cuda()\n",
    "\n",
    "for i in latent_dim_lst:\n",
    "    print('bottleneck dimension: {}'.format(i))\n",
    "    model = fm.load_model('{}_{}'.format(modelname, i)).to(device)\n",
    "    output = model(img.to(device))\n",
    "\n",
    "    original = pl.to_img(img.cpu().data, [3, 218, 178]) \n",
    "    pic = pl.to_img(output.cpu().data, [3, 218, 178])\n",
    "\n",
    "    print(\"Original Image:\")\n",
    "    pl.imshow(torchvision.utils.make_grid(original[:num_img].detach(), grid_row_size), filename=\"com_classic_celeba_{}_original\".format(i))\n",
    "    print(\"Reconstructed Image:\")\n",
    "    pl.imshow(torchvision.utils.make_grid(pic[:num_img].detach(), grid_row_size), filename=\"com_classic_celeba_{}_reconstructed\".format(i))\n",
    "    print(\"Difference:\")\n",
    "    diff_img = (original - pic + 1) / 2\n",
    "    pl.imshow(torchvision.utils.make_grid(diff_img[:num_img].detach(), grid_row_size), filename=\"com_classic_celeba_{}_difference\".format(i))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot Recontruction Loss against Bottleneck Size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = fm.load_variable(\"{}_bottleneck\".format(modelname))\n",
    "y = [train, test]\n",
    "\n",
    "pl.plot(latent_dim_lst, y, 'bottleneck size', 'loss', ['train', 'test'], 'Train & Test Reconstruction Loss History', 'loss_l1_celebA_bottleneck') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}