pytorch/Adversarial Autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn\n",
    "import torch.nn.functional as nn\n",
    "import torch.autograd as autograd\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import os\n",
    "from torch.autograd import Variable\n",
    "from tensorflow.examples.tutorials.mnist import input_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully downloaded train-images-idx3-ubyte.gz 9912422 bytes.\n",
      "Extracting ../../MNIST_data/train-images-idx3-ubyte.gz\n",
      "Successfully downloaded train-labels-idx1-ubyte.gz 28881 bytes.\n",
      "Extracting ../../MNIST_data/train-labels-idx1-ubyte.gz\n",
      "Successfully downloaded t10k-images-idx3-ubyte.gz 1648877 bytes.\n",
      "Extracting ../../MNIST_data/t10k-images-idx3-ubyte.gz\n",
      "Successfully downloaded t10k-labels-idx1-ubyte.gz 4542 bytes.\n",
      "Extracting ../../MNIST_data/t10k-labels-idx1-ubyte.gz\n"
     ]
    }
   ],
   "source": [
    "mnist = input_data.read_data_sets('../../MNIST_data', one_hot=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mb_size = 32\n",
    "z_dim = 5\n",
    "X_dim = mnist.train.images.shape[1]\n",
    "y_dim = mnist.train.labels.shape[1]\n",
    "h_dim = 128\n",
    "cnt = 0\n",
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Encoder\n",
    "Q = torch.nn.Sequential(\n",
    "    torch.nn.Linear(X_dim, h_dim),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(h_dim, z_dim)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decoder\n",
    "P = torch.nn.Sequential(\n",
    "    torch.nn.Linear(z_dim, h_dim),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(h_dim, X_dim),\n",
    "    torch.nn.Sigmoid())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Discriminator\n",
    "D = torch.nn.Sequential(\n",
    "    torch.nn.Linear(z_dim, h_dim),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(h_dim, 1),\n",
    "    torch.nn.Sigmoid()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Reset Gradient\n",
    "def reset_grad():\n",
    "    Q.zero_grad()\n",
    "    P.zero_grad()\n",
    "    D.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sample_X(size, include_y=False):\n",
    "    X, y = mnist.train.next_batch(size)\n",
    "    X = Variable(torch.from_numpy(X))\n",
    "\n",
    "    if include_y:\n",
    "        y = np.argmax(y, axis=1).astype(np.int)\n",
    "        y = Variable(torch.from_numpy(y))\n",
    "        return X, y\n",
    "\n",
    "    return X\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter-0; D_loss: 1.439; G_loss: 0.7287; recon_loss: 0.6982\n",
      "Iter-1000; D_loss: 1.371; G_loss: 0.7533; recon_loss: 0.2727\n",
      "Iter-2000; D_loss: 1.449; G_loss: 0.6482; recon_loss: 0.2138\n",
      "Iter-3000; D_loss: 1.408; G_loss: 0.69; recon_loss: 0.1909\n",
      "Iter-4000; D_loss: 1.383; G_loss: 0.7024; recon_loss: 0.1826\n",
      "Iter-5000; D_loss: 1.39; G_loss: 0.6931; recon_loss: 0.177\n",
      "Iter-6000; D_loss: 1.372; G_loss: 0.7667; recon_loss: 0.1709\n",
      "Iter-7000; D_loss: 1.395; G_loss: 0.7129; recon_loss: 0.1825\n",
      "Iter-8000; D_loss: 1.389; G_loss: 0.6816; recon_loss: 0.1665\n",
      "Iter-9000; D_loss: 1.39; G_loss: 0.6768; recon_loss: 0.1914\n",
      "Iter-10000; D_loss: 1.387; G_loss: 0.6906; recon_loss: 0.1478\n",
      "Iter-11000; D_loss: 1.379; G_loss: 0.7249; recon_loss: 0.167\n",
      "Iter-12000; D_loss: 1.393; G_loss: 0.6823; recon_loss: 0.1833\n",
      "Iter-13000; D_loss: 1.386; G_loss: 0.6821; recon_loss: 0.1486\n",
      "Iter-14000; D_loss: 1.393; G_loss: 0.6952; recon_loss: 0.1572\n",
      "Iter-15000; D_loss: 1.386; G_loss: 0.7; recon_loss: 0.1638\n",
      "Iter-16000; D_loss: 1.383; G_loss: 0.696; recon_loss: 0.1668\n",
      "Iter-17000; D_loss: 1.391; G_loss: 0.6997; recon_loss: 0.163\n",
      "Iter-18000; D_loss: 1.388; G_loss: 0.6924; recon_loss: 0.1619\n",
      "Iter-19000; D_loss: 1.388; G_loss: 0.6861; recon_loss: 0.1596\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-35e7d869027e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0mG_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m     \u001b[0mQ_solver\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     42\u001b[0m     \u001b[0mreset_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda/lib/python3.6/site-packages/torch/optim/adam.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m     72\u001b[0m                 \u001b[0mstep_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgroup\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'lr'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mmath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbias_correction2\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mbias_correction1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     73\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 74\u001b[0;31m                 \u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maddcdiv_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mstep_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexp_avg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdenom\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     75\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     76\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "Q_solver = optim.Adam(Q.parameters(), lr=lr)\n",
    "P_solver = optim.Adam(P.parameters(), lr=lr)\n",
    "D_solver = optim.Adam(D.parameters(), lr=lr)\n",
    "\n",
    "\n",
    "for it in range(1000000):\n",
    "    X = sample_X(mb_size)\n",
    "\n",
    "    \"\"\" Reconstruction phase \"\"\"\n",
    "    z_sample = Q(X)\n",
    "    X_sample = P(z_sample)\n",
    "\n",
    "    recon_loss = nn.binary_cross_entropy(X_sample, X)\n",
    "\n",
    "    recon_loss.backward()\n",
    "    P_solver.step()\n",
    "    Q_solver.step()\n",
    "    reset_grad()\n",
    "\n",
    "    \"\"\" Regularization phase \"\"\"\n",
    "    # Discriminator\n",
    "    z_real = Variable(torch.randn(mb_size, z_dim))\n",
    "    z_fake = Q(X)\n",
    "\n",
    "    D_real = D(z_real)\n",
    "    D_fake = D(z_fake)\n",
    "\n",
    "    D_loss = -torch.mean(torch.log(D_real) + torch.log(1 - D_fake))\n",
    "\n",
    "    D_loss.backward()\n",
    "    D_solver.step()\n",
    "    reset_grad()\n",
    "\n",
    "    # Generator\n",
    "    z_fake = Q(X)\n",
    "    D_fake = D(z_fake)\n",
    "\n",
    "    G_loss = -torch.mean(torch.log(D_fake))\n",
    "\n",
    "    G_loss.backward()\n",
    "    Q_solver.step()\n",
    "    reset_grad()\n",
    "\n",
    "    # Print and plot every now and then\n",
    "    if it % 1000 == 0:\n",
    "        print('Iter-{}; D_loss: {:.4}; G_loss: {:.4}; recon_loss: {:.4}'\n",
    "              .format(it, D_loss.data[0], G_loss.data[0], recon_loss.data[0]))\n",
    "\n",
    "        samples = P(z_real).data.numpy()[:16]\n",
    "\n",
    "        fig = plt.figure(figsize=(4, 4))\n",
    "        gs = gridspec.GridSpec(4, 4)\n",
    "        gs.update(wspace=0.05, hspace=0.05)\n",
    "\n",
    "        for i, sample in enumerate(samples):\n",
    "            ax = plt.subplot(gs[i])\n",
    "            plt.axis('off')\n",
    "            ax.set_xticklabels([])\n",
    "            ax.set_yticklabels([])\n",
    "            ax.set_aspect('equal')\n",
    "            plt.imshow(sample.reshape(28, 28), cmap='Greys_r')\n",
    "\n",
    "        if not os.path.exists('out/'):\n",
    "            os.makedirs('out/')\n",
    "\n",
    "        plt.savefig('out/{}.png'\n",
    "                    .format(str(cnt).zfill(3)), bbox_inches='tight')\n",
    "        cnt += 1\n",
    "        plt.close(fig)"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}