VAE-GAN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import argparse\n",
    "import h5py\n",
    "import numpy as np\n",
    "import os\n",
    "import time\n",
    "import torch\n",
    "import torch.utils.data\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import make_grid , save_image\n",
    "import torchvision.utils as vutils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def show_and_save(file_name,img):\n",
    "    npimg = np.transpose(img.numpy(),(1,2,0))\n",
    "    f = \"./%s.png\" % file_name\n",
    "    fig = plt.figure(dpi=300)\n",
    "    fig.suptitle(file_name, fontsize=14, fontweight='bold')\n",
    "    #plt.imshow(npimg)\n",
    "    plt.imsave(f,npimg)\n",
    "    \n",
    "def save_model(epoch, encoder, decoder, D):\n",
    "    torch.save(decoder.cpu().state_dict(), './VAE_GAN_decoder_%d.pth' % epoch)\n",
    "    torch.save(encoder.cpu().state_dict(),'./VAE_GAN_encoder_%d.pth' % epoch)\n",
    "    torch.save(D.cpu().state_dict(), 'VAE_GAN_D_%d.pth' % epoch)\n",
    "    decoder.cuda()\n",
    "    encoder.cuda()\n",
    "    D.cuda()\n",
    "    \n",
    "def load_model(epoch, encoder, decoder, D):\n",
    "    #  restore models\n",
    "    decoder.load_state_dict(torch.load('./VAE_GAN_decoder_%d.pth' % epoch))\n",
    "    decoder.cuda()\n",
    "    encoder.load_state_dict(torch.load('./VAE_GAN_encoder_%d.pth' % epoch))\n",
    "    encoder.cuda()\n",
    "    D.load_state_dict(torch.load('VAE_GAN_D_%d.pth' % epoch))\n",
    "    D.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CelebADataset(Dataset):\n",
    "    def __init__(self, h5_path, transform=None):\n",
    "        assert (os.path.isfile(h5_path))\n",
    "        self.h5_path = h5_path\n",
    "        self.transform = transform\n",
    "        \n",
    "        # loading the dataset into memory\n",
    "        f = h5py.File(self.h5_path, \"r\")\n",
    "        key = list(f.keys())\n",
    "        print (\"key list:\", key)\n",
    "        self.dataset = f[key[0]]\n",
    "        self.dataset = self.dataset[:1000]\n",
    "        print (\"dataset loaded and its shape:\", self.dataset.shape)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        img = self.dataset[index]\n",
    "        img = np.transpose(img, (1, 2, 0))\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "            \n",
    "        return img, 0\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "T = transforms.Compose([transforms.ToTensor(),\n",
    "                        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),])\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    CelebADataset('../dataset/CelebA_aligned_reduced.h5', transform=T),\n",
    "    batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data, _ = next(iter(train_loader))\n",
    "print (data.size())\n",
    "print (data[1, :])\n",
    "show_and_save(\"gen\" ,make_grid((data*0.5+0.5).cpu(),8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_channels, output_channels, representation_size = 64):\n",
    "        super(Encoder, self).__init__()\n",
    "        # input parameters\n",
    "        self.input_channels = input_channels\n",
    "        self.output_channels = output_channels\n",
    "        \n",
    "        self.features = nn.Sequential(\n",
    "            # nc x 64 x 64\n",
    "            nn.Conv2d(self.input_channels, representation_size, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(representation_size),\n",
    "            nn.ReLU(),\n",
    "            # hidden_size x 32 x 32\n",
    "            nn.Conv2d(representation_size, representation_size*2, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(representation_size * 2),\n",
    "            nn.ReLU(),\n",
    "            # hidden_size*2 x 16 x 16\n",
    "            nn.Conv2d(representation_size*2, representation_size*4, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(representation_size * 4),\n",
    "            nn.ReLU())\n",
    "            # hidden_size*4 x 8 x 8\n",
    "            \n",
    "        self.mean = nn.Sequential(\n",
    "            nn.Linear(representation_size*4*8*8, 2048),\n",
    "            nn.BatchNorm1d(2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(2048, output_channels))\n",
    "        \n",
    "        self.logvar = nn.Sequential(\n",
    "            nn.Linear(representation_size*4*8*8, 2048),\n",
    "            nn.BatchNorm1d(2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(2048, output_channels))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "\n",
    "        hidden_representation = self.features(x)\n",
    "\n",
    "        mean = self.mean(hidden_representation.view(batch_size, -1))\n",
    "        logvar = self.logvar(hidden_representation.view(batch_size, -1))\n",
    "\n",
    "        return mean, logvar\n",
    "    \n",
    "    def hidden_layer(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "        output = self.features(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, input_size, representation_size):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.representation_size = representation_size\n",
    "        dim = representation_size[0] * representation_size[1] * representation_size[2]\n",
    "        \n",
    "        self.preprocess = nn.Sequential(\n",
    "            nn.Linear(input_size, dim),\n",
    "            nn.BatchNorm1d(dim),\n",
    "            nn.ReLU())\n",
    "        \n",
    "            # 256 x 8 x 8\n",
    "        self.deconv1 = nn.ConvTranspose2d(representation_size[0], 256, 5, stride=2, padding=2)\n",
    "        self.act1 = nn.Sequential(nn.BatchNorm2d(256),\n",
    "                                  nn.ReLU())\n",
    "            # 256 x 16 x 16\n",
    "        self.deconv2 = nn.ConvTranspose2d(256, 128, 5, stride=2, padding=2)\n",
    "        self.act2 = nn.Sequential(nn.BatchNorm2d(128),\n",
    "                                  nn.ReLU())\n",
    "            # 128 x 32 x 32\n",
    "        self.deconv3 = nn.ConvTranspose2d(128, 32, 5, stride=2, padding=2)\n",
    "        self.act3 = nn.Sequential(nn.BatchNorm2d(32),\n",
    "                                  nn.ReLU())\n",
    "            # 32 x 64 x 64\n",
    "        self.deconv4 = nn.ConvTranspose2d(32, 3, 5, stride=1, padding=2)\n",
    "            # 3 x 64 x 64\n",
    "        self.activation = nn.Tanh()\n",
    "            \n",
    "    \n",
    "    def forward(self, code):\n",
    "        bs = code.size()[0]\n",
    "        preprocessed_codes = self.preprocess(code)\n",
    "        preprocessed_codes = preprocessed_codes.view(-1,\n",
    "                                                     self.representation_size[0],\n",
    "                                                     self.representation_size[1],\n",
    "                                                     self.representation_size[2])\n",
    "        output = self.deconv1(preprocessed_codes, output_size=(bs, 256, 16, 16))\n",
    "        output = self.act1(output)\n",
    "        output = self.deconv2(output, output_size=(bs, 128, 32, 32))\n",
    "        output = self.act2(output)\n",
    "        output = self.deconv3(output, output_size=(bs, 32, 64, 64))\n",
    "        output = self.act3(output)\n",
    "        output = self.deconv4(output, output_size=(bs, 3, 64, 64))\n",
    "        output = self.activation(output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE_GAN_Generator(nn.Module):\n",
    "    def __init__(self, input_channels, hidden_size, representation_size=(256, 8, 8)):\n",
    "        super(VAE_GAN_Generator, self).__init__()\n",
    "        self.input_channels = input_channels\n",
    "        self.hidden_size = hidden_size\n",
    "        self.representation_size = representation_size\n",
    "        \n",
    "        self.encoder = Encoder(input_channels, hidden_size)\n",
    "        self.decoder = Decoder(hidden_size, representation_size)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "        mean, logvar = self.encoder(x)\n",
    "        std = logvar.mul(0.5).exp_()\n",
    "        \n",
    "        reparametrized_noise = Variable(torch.randn((batch_size, self.hidden_size))).cuda()\n",
    "\n",
    "        reparametrized_noise = mean + std * reparametrized_noise\n",
    "\n",
    "        rec_images = self.decoder(reparametrized_noise)\n",
    "        \n",
    "        return mean, logvar, rec_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_channels, representation_size=(256, 8, 8)):  \n",
    "        super(Discriminator, self).__init__()\n",
    "        self.representation_size = representation_size\n",
    "        dim = representation_size[0] * representation_size[1] * representation_size[2]\n",
    "        \n",
    "        self.main = nn.Sequential(\n",
    "            nn.Conv2d(input_channels, 32, 5, stride=1, padding=2),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Conv2d(32, 128, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Conv2d(128, 256, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Conv2d(256, 256, 5, stride=2, padding=2),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2))\n",
    "        \n",
    "        self.lth_features = nn.Sequential(\n",
    "            nn.Linear(dim, 2048),\n",
    "            nn.LeakyReLU(0.2))\n",
    "        \n",
    "        self.sigmoid_output = nn.Sequential(\n",
    "            nn.Linear(2048, 1),\n",
    "            nn.Sigmoid())\n",
    "        \n",
    "    def forward(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "        features = self.main(x)\n",
    "        lth_rep = self.lth_features(features.view(batch_size, -1))\n",
    "        output = self.sigmoid_output(lth_rep)\n",
    "        return output\n",
    "    \n",
    "    def similarity(self, x):\n",
    "        batch_size = x.size()[0]\n",
    "        features = self.main(x)\n",
    "        lth_rep = self.lth_features(features.view(batch_size, -1))\n",
    "        return lth_rep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define constant\n",
    "input_channels = 3\n",
    "hidden_size = 64\n",
    "max_epochs = 250\n",
    "lr = 3e-4\n",
    "\n",
    "beta = 5\n",
    "alpha = 0.1\n",
    "gamma = 15\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = VAE_GAN_Generator(input_channels, hidden_size).cuda()\n",
    "D = Discriminator(input_channels).cuda()\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "criterion.cuda()\n",
    "\n",
    "opt_enc = optim.RMSprop(G.encoder.parameters(), lr=lr)\n",
    "opt_dec = optim.RMSprop(G.decoder.parameters(), lr=lr)\n",
    "opt_dis = optim.RMSprop(D.parameters(), lr=lr * alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fixed_noise = Variable(torch.randn(batch_size, hidden_size)).cuda()\n",
    "data, _ = next(iter(train_loader))\n",
    "fixed_batch = Variable(data).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for epoch in range(max_epochs):\n",
    "    D_real_list, D_rec_enc_list, D_rec_noise_list, D_list = [], [], [], []\n",
    "    g_loss_list, rec_loss_list, prior_loss_list = [], [], []\n",
    "    for data, _ in train_loader:\n",
    "        batch_size = data.size()[0]\n",
    "        ones_label = Variable(torch.ones(batch_size)).cuda()\n",
    "        zeros_label = Variable(torch.zeros(batch_size)).cuda()\n",
    "        \n",
    "        #print (data.size())\n",
    "        datav = Variable(data).cuda()\n",
    "        mean, logvar, rec_enc = G(datav)\n",
    "        #print (\"The size of rec_enc:\", rec_enc.size())\n",
    "        \n",
    "        noisev = Variable(torch.randn(batch_size, hidden_size)).cuda()\n",
    "        rec_noise = G.decoder(noisev)\n",
    "        \n",
    "        # train discriminator\n",
    "        output = D(datav)\n",
    "        errD_real = criterion(output, ones_label)\n",
    "        D_real_list.append(output.data.mean())\n",
    "        output = D(rec_enc)\n",
    "        errD_rec_enc = criterion(output, zeros_label)\n",
    "        D_rec_enc_list.append(output.data.mean())\n",
    "        output = D(rec_noise)\n",
    "        errD_rec_noise = criterion(output, zeros_label)\n",
    "        D_rec_noise_list.append(output.data.mean())\n",
    "        \n",
    "        dis_img_loss = errD_real + errD_rec_enc + errD_rec_noise\n",
    "        #print (\"print (dis_img_loss)\", dis_img_loss)\n",
    "        D_list.append(dis_img_loss.data.mean())\n",
    "        opt_dis.zero_grad()\n",
    "        dis_img_loss.backward(retain_graph=True)\n",
    "        opt_dis.step()\n",
    "        \n",
    "        # train decoder\n",
    "        output = D(datav)\n",
    "        errD_real = criterion(output, ones_label)\n",
    "        output = D(rec_enc)\n",
    "        errD_rec_enc = criterion(output, zeros_label)\n",
    "        output = D(rec_noise)\n",
    "        errD_rec_noise = criterion(output, zeros_label)\n",
    "        \n",
    "        similarity_rec_enc = D.similarity(rec_enc)\n",
    "        similarity_data = D.similarity(datav)\n",
    "        \n",
    "        dis_img_loss = errD_real + errD_rec_enc + errD_rec_noise\n",
    "        #print (dis_img_loss)\n",
    "        gen_img_loss = - dis_img_loss\n",
    "        \n",
    "        g_loss_list.append(gen_img_loss.data.mean())\n",
    "        rec_loss = ((similarity_rec_enc - similarity_data) ** 2).mean()\n",
    "        rec_loss_list.append(rec_loss.data.mean())\n",
    "        err_dec = gamma * rec_loss + gen_img_loss\n",
    "        \n",
    "        opt_dec.zero_grad()\n",
    "        err_dec.backward(retain_graph=True)\n",
    "        opt_dec.step()\n",
    "        \n",
    "        # train encoder\n",
    "        prior_loss = 1 + logvar - mean.pow(2) - logvar.exp()\n",
    "        prior_loss = (-0.5 * torch.sum(prior_loss))/torch.numel(mean.data)\n",
    "        #print (prior_loss, mean, std)\n",
    "        prior_loss_list.append(prior_loss.data.mean())\n",
    "        err_enc = prior_loss + beta * rec_loss\n",
    "        \n",
    "        opt_enc.zero_grad()\n",
    "        err_enc.backward()\n",
    "        opt_enc.step()\n",
    "        \n",
    "    #save_model(epoch, G.encoder, G.decoder, D)\n",
    "    _, _, rec_imgs = G(fixed_batch)\n",
    "    show_and_save('rec_epoch_%d.png' % epoch ,make_grid((rec_imgs.data*0.5+0.5).cpu(),8))\n",
    "    '''\n",
    "    vutils.save_image(rec_imgs.data,\n",
    "            'rec_epoch_%d.png' % epoch,\n",
    "            normalize=True)\n",
    "    '''\n",
    "    samples = G.decoder(fixed_noise)\n",
    "    vutils.save_image(samples.data,\n",
    "            'sample_epoch_%d.png' % epoch,\n",
    "            normalize=True)\n",
    "    localtime = time.asctime( time.localtime(time.time()) )\n",
    "    print (localtime)\n",
    "    print ('[%d/%d]: D_real:%.4f, D_enc:%.4f, D_noise:%.4f, Loss_D:%.4f, Loss_G:%.4f, rec_loss:%.4f, prior_loss:%.4f' \n",
    "           % (epoch, \n",
    "              max_epochs, \n",
    "              np.mean(D_real_list), \n",
    "              np.mean(D_rec_enc_list), \n",
    "              np.mean(D_rec_noise_list), \n",
    "              np.mean(D_list), \n",
    "              np.mean(g_loss_list),\n",
    "              np.mean(rec_loss_list),\n",
    "              np.mean(prior_loss_list)))\n",
    "    "
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}