vae_script_celeba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8bcd9f4",
   "metadata": {},
   "source": [
    "# VAE CelebA\n",
    "### Overview\n",
    "This script is used to train a VAE (BetaVAE,  𝛽=1 ), BetaVAE, FactorVAE, or BetaTCVAE on the CelebA dataset, then save it. The script will train a VAE-based model on a fixed amount of data using the hyperparameters defined in the cell below. The script will train the network with the given hyperparameters, compare original data with reconstructions, create images of latent traversals, and save the model.\n",
    "\n",
    "### Instructions\n",
    "Set hyperparameters for the run in the cell below. Then, hit Run All on the jupyter notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fe16cbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### CHOOSE A SEED #### (or generate randomly)\n",
    "seed = 60\n",
    "import random\n",
    "random.seed(seed)\n",
    "\n",
    "from ae_utils_exp import B_TCVAE as VAE_BASED_MODEL # change <model> in \".... import <model> as ....\"\n",
    "### options: VAE (for VAE, BetaVAE), FACTOR_VAE, or B_TCVAE (for BetaTCVAE)\n",
    "\n",
    "### SELECT HYPERPARAMETERS FOR THIS RUN #######\n",
    "\n",
    "beta = 50.0 # \\beta for BetaVae, FactorVAE, and BetaTCVAE# \\beta for BetaVae, FactorVAE, and BetaTCVAE\n",
    "n_lat = 32 # VAE bottleneck size (m) \n",
    "batch_size = 200 # batch size used for training\n",
    "lr = 0.001 # learning rate used for training. should be 1e-4 if the model is FACTOR_VAE\n",
    "\n",
    "savename = \"./models/celeba_vae_lr1e-3_seed{}_b{}.pt\".format(seed, beta) # savename for the trained model\n",
    "\n",
    "\n",
    "print(\"Seed: \", seed)\n",
    "print(\"Beta: \", beta)\n",
    "print(\"Batch Size: \", batch_size)\n",
    "print(\"Savename: \", savename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e909c01a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ae_utils_exp import multi_t\n",
    "from torchvision.transforms import Compose\n",
    "\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea73cbbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00559994",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.datasets import CelebA\n",
    "import torchvision.transforms as tforms\n",
    "\n",
    "tform = tforms.Compose([tforms.Resize(96), tforms.CenterCrop(64), tforms.ToTensor()])\n",
    "\n",
    "dataset = CelebA(root='../beamsynthesizer/data', split='all', download=False, transform=tform)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0327f7c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ae_utils_exp import celeba_norm, celeba_inorm\n",
    "from architectures import enc_celeba_small_vae as enc\n",
    "from architectures import dec_celeba_small as dec\n",
    "\n",
    "ae = VAE_BASED_MODEL(celeba_norm, enc(lat=n_lat, inp_chan=3), dec(lat=n_lat, inp_chan=3), \\\n",
    "                     device, z_dim=n_lat, inp_inorm=celeba_inorm)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "170fb753",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_loss, kl_loss = \\\n",
    "    ae.fit(dataset, 200, batch_per_group=20, batch_size=batch_size, lr=lr, beta=beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf7f9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n",
    "# plot the loss curves on a log scale\n",
    "ax[0].set_ylabel(\"$log_{10}$(LogProb Rec Loss)\")\n",
    "ax[0].set_xlabel(\"Group\")\n",
    "ax[0].plot(np.log10(rec_loss), linewidth=2, label='Reconstruction')\n",
    "ax[0].legend()\n",
    "ax[0].grid(True, which='both', ls='-')\n",
    "\n",
    "ax[1].set_ylabel(\"$log_{10}$(KL Loss)\")\n",
    "ax[1].set_xlabel(\"Group\")\n",
    "ax[1].plot(np.log10(kl_loss), linewidth=2, label='KL')\n",
    "ax[1].legend()\n",
    "ax[1].grid(True, which='both', ls='-')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e81e872",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_batch_size=200\n",
    "num_to_plot=20\n",
    "z_scores, z_pred_scores, inp, rec = ae.record_latent_space(dataset, batch_size=plt_batch_size, n_batches=5)\n",
    "\n",
    "inp = multi_t(inp, 1, 3).clamp(0, 1).cpu().numpy()\n",
    "rec = multi_t(rec, 1, 3).clamp(0, 1).cpu().numpy()\n",
    "\n",
    "fig, axes = plt.subplots(2, num_to_plot, figsize=(20, 4))\n",
    "for i in range(num_to_plot):\n",
    "    axes[0][i].imshow(inp[i])\n",
    "    axes[1][i].imshow(rec[i])\n",
    "    axes[0][i].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n",
    "    axes[1][i].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06883920",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ae_utils_exp import InvNorm\n",
    "\n",
    "invn = celeba_inorm\n",
    "\n",
    "\n",
    "# determine base z_scores\n",
    "ind = 0\n",
    "z_base = z_scores[ind]\n",
    "fig, axes = plt.subplots(1, 2, figsize=(6, 3))\n",
    "axes[0].imshow(inp[ind], cmap='gray')\n",
    "axes[0].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n",
    "axes[1].imshow(rec[ind], cmap='gray')\n",
    "axes[1].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f92be45c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# decode\n",
    "fig, axes = plt.subplots(ae.z_dim//2, 10, figsize=(16, 24))\n",
    "with torch.no_grad():\n",
    "    for i in range(ae.z_dim//2):\n",
    "        _min = z_scores[:, i].min()\n",
    "        _max = z_scores[:, i].max()\n",
    "        variation = torch.linspace(_min, _max, steps=10)\n",
    "        for j in range(len(variation)):\n",
    "            axes[i][j].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n",
    "            if _max - _min >= 0.2:\n",
    "                z = z_base.clone()\n",
    "                z[i] = variation[j]\n",
    "                im = multi_t(invn(ae.dec(z.to(ae.device))), 1, 3).clamp(0, 1).squeeze().cpu().numpy()\n",
    "                axes[i][j].imshow(im, cmap='gray')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c184c093",
   "metadata": {},
   "outputs": [],
   "source": [
    "# decode\n",
    "fig, axes = plt.subplots(ae.z_dim//2, 10, figsize=(16, 24))\n",
    "with torch.no_grad():\n",
    "    for i in range(ae.z_dim//2):\n",
    "        _min = z_scores[:, i + ae.z_dim//2].min()\n",
    "        _max = z_scores[:, i + ae.z_dim//2].max()\n",
    "        variation = torch.linspace(_min, _max, steps=10)\n",
    "        for j in range(len(variation)):\n",
    "            axes[i][j].tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, labelleft=False)\n",
    "            if _max - _min >= 0.2:\n",
    "                z = z_base.clone()\n",
    "                z[i + ae.z_dim//2] = variation[j]\n",
    "                im = multi_t(invn(ae.dec(z.to(ae.device))), 1, 3).clamp(0,1).squeeze().cpu().numpy()\n",
    "                axes[i][j].imshow(im, cmap='gray')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d741b9d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(ae.state_dict(), savename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e1e286",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}