WGAN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting MNIST_data/train-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/train-labels-idx1-ubyte.gz\n",
      "Extracting MNIST_data/t10k-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/t10k-labels-idx1-ubyte.gz\n"
     ]
    }
   ],
   "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\n",
    "\n",
    "\n",
    "mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)\n",
    "X_dim = mnist.train.images.shape[1]\n",
    "y_dim = mnist.train.labels.shape[1]\n",
    "h_dim = 128\n",
    "z_dim = 10\n",
    "cnt = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = 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()\n",
    ")\n",
    "\n",
    "\n",
    "D = torch.nn.Sequential(\n",
    "    torch.nn.Linear(X_dim, h_dim),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(h_dim, 1),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def reset_grad():\n",
    "    G.zero_grad()\n",
    "    D.zero_grad()\n",
    "\n",
    "\n",
    "G_solver = optim.RMSprop(G.parameters(), lr=0.0001)\n",
    "D_solver = optim.RMSprop(D.parameters(), lr=0.0001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter-0; D_loss: [-0.12624831]; G_loss: [ 0.16157164]\n",
      "Iter-1000; D_loss: [-0.03480838]; G_loss: [-0.00120416]\n",
      "Iter-2000; D_loss: [-0.00738919]; G_loss: [ 0.0535388]\n",
      "Iter-3000; D_loss: [-0.05886087]; G_loss: [ 0.00255544]\n",
      "Iter-4000; D_loss: [-0.04449232]; G_loss: [ 0.00304415]\n",
      "Iter-5000; D_loss: [-0.03686142]; G_loss: [-0.00836258]\n",
      "Iter-6000; D_loss: [-0.03010887]; G_loss: [-0.00741883]\n",
      "Iter-7000; D_loss: [-0.03416763]; G_loss: [-0.0095391]\n",
      "Iter-8000; D_loss: [-0.03482541]; G_loss: [-0.00438419]\n",
      "Iter-9000; D_loss: [-0.0277044]; G_loss: [-0.0063392]\n",
      "Iter-10000; D_loss: [-0.02745331]; G_loss: [-0.00485104]\n",
      "Iter-11000; D_loss: [-0.0272183]; G_loss: [-0.00546759]\n",
      "Iter-12000; D_loss: [-0.02287278]; G_loss: [-0.00616492]\n",
      "Iter-13000; D_loss: [-0.02330181]; G_loss: [-0.00553512]\n",
      "Iter-14000; D_loss: [-0.02414522]; G_loss: [-0.00600764]\n",
      "Iter-15000; D_loss: [-0.02072017]; G_loss: [-0.00366222]\n",
      "Iter-16000; D_loss: [-0.02135508]; G_loss: [-0.00627525]\n",
      "Iter-17000; D_loss: [-0.01786141]; G_loss: [-0.00336719]\n",
      "Iter-18000; D_loss: [-0.02380688]; G_loss: [-0.00715477]\n",
      "Iter-19000; D_loss: [-0.01739286]; G_loss: [-0.00733081]\n",
      "Iter-20000; D_loss: [-0.01940097]; G_loss: [-0.0044988]\n",
      "Iter-21000; D_loss: [-0.01271411]; G_loss: [-0.00444629]\n",
      "Iter-22000; D_loss: [-0.01605458]; G_loss: [-0.00367484]\n",
      "Iter-23000; D_loss: [-0.01635769]; G_loss: [-0.00460203]\n",
      "Iter-24000; D_loss: [-0.01213492]; G_loss: [-0.00709497]\n",
      "Iter-25000; D_loss: [-0.01760268]; G_loss: [-0.00774769]\n",
      "Iter-26000; D_loss: [-0.01409053]; G_loss: [-0.00909518]\n",
      "Iter-27000; D_loss: [-0.01406839]; G_loss: [-0.00483367]\n",
      "Iter-28000; D_loss: [-0.01397273]; G_loss: [-0.00725979]\n",
      "Iter-29000; D_loss: [-0.01267296]; G_loss: [-0.00793647]\n",
      "Iter-30000; D_loss: [-0.01413301]; G_loss: [-0.01077011]\n",
      "Iter-31000; D_loss: [-0.01466979]; G_loss: [-0.00912927]\n",
      "Iter-32000; D_loss: [-0.01465054]; G_loss: [-0.00971281]\n",
      "Iter-33000; D_loss: [-0.01213105]; G_loss: [-0.00684607]\n",
      "Iter-34000; D_loss: [-0.01093945]; G_loss: [-0.00665584]\n",
      "Iter-35000; D_loss: [-0.01356145]; G_loss: [-0.00735524]\n",
      "Iter-36000; D_loss: [-0.01252815]; G_loss: [-0.01235853]\n",
      "Iter-37000; D_loss: [-0.01224817]; G_loss: [-0.00677556]\n",
      "Iter-38000; D_loss: [-0.01192078]; G_loss: [-0.00873478]\n",
      "Iter-39000; D_loss: [-0.01211099]; G_loss: [-0.00865818]\n",
      "Iter-40000; D_loss: [-0.01443424]; G_loss: [-0.00933674]\n",
      "Iter-41000; D_loss: [-0.01346017]; G_loss: [-0.01021007]\n",
      "Iter-42000; D_loss: [-0.01234387]; G_loss: [-0.01217627]\n",
      "Iter-43000; D_loss: [-0.01089216]; G_loss: [-0.00632808]\n",
      "Iter-44000; D_loss: [-0.01035895]; G_loss: [-0.00725813]\n",
      "Iter-45000; D_loss: [-0.01147261]; G_loss: [-0.01133318]\n",
      "Iter-46000; D_loss: [-0.01488632]; G_loss: [-0.01060204]\n",
      "Iter-47000; D_loss: [-0.01173247]; G_loss: [-0.01153311]\n",
      "Iter-48000; D_loss: [-0.00731252]; G_loss: [-0.00851685]\n",
      "Iter-49000; D_loss: [-0.01288038]; G_loss: [-0.01170281]\n",
      "Iter-50000; D_loss: [-0.00837466]; G_loss: [-0.01010827]\n",
      "Iter-51000; D_loss: [-0.01116705]; G_loss: [-0.00839094]\n",
      "Iter-52000; D_loss: [-0.01062771]; G_loss: [-0.01208306]\n",
      "Iter-53000; D_loss: [-0.01118516]; G_loss: [-0.01028057]\n",
      "Iter-54000; D_loss: [-0.00954585]; G_loss: [-0.00859939]\n",
      "Iter-55000; D_loss: [-0.00977315]; G_loss: [-0.00814354]\n",
      "Iter-56000; D_loss: [-0.01107703]; G_loss: [-0.00686317]\n",
      "Iter-57000; D_loss: [-0.00947806]; G_loss: [-0.0094428]\n",
      "Iter-58000; D_loss: [-0.00674607]; G_loss: [-0.00948315]\n",
      "Iter-59000; D_loss: [-0.00878486]; G_loss: [-0.00781303]\n",
      "Iter-60000; D_loss: [-0.00894366]; G_loss: [-0.00852806]\n",
      "Iter-61000; D_loss: [-0.00843244]; G_loss: [-0.00952837]\n",
      "Iter-62000; D_loss: [-0.0091804]; G_loss: [-0.00568749]\n",
      "Iter-63000; D_loss: [-0.01158731]; G_loss: [-0.00374322]\n",
      "Iter-64000; D_loss: [-0.01274681]; G_loss: [-0.00415864]\n",
      "Iter-65000; D_loss: [-0.01044841]; G_loss: [-0.00634096]\n",
      "Iter-66000; D_loss: [-0.00860463]; G_loss: [-0.00986252]\n",
      "Iter-67000; D_loss: [-0.01083128]; G_loss: [-0.00727835]\n",
      "Iter-68000; D_loss: [-0.00649342]; G_loss: [-0.00871784]\n",
      "Iter-69000; D_loss: [-0.00781219]; G_loss: [-0.00719225]\n",
      "Iter-70000; D_loss: [-0.00815525]; G_loss: [-0.00727612]\n",
      "Iter-71000; D_loss: [-0.01122109]; G_loss: [-0.01047037]\n",
      "Iter-72000; D_loss: [-0.01030962]; G_loss: [-0.00740754]\n",
      "Iter-73000; D_loss: [-0.01089522]; G_loss: [-0.0063086]\n",
      "Iter-74000; D_loss: [-0.01018651]; G_loss: [-0.00629523]\n",
      "Iter-75000; D_loss: [-0.00801003]; G_loss: [-0.00794108]\n",
      "Iter-76000; D_loss: [-0.01071993]; G_loss: [-0.00425647]\n",
      "Iter-77000; D_loss: [-0.00706098]; G_loss: [-0.00541867]\n",
      "Iter-78000; D_loss: [-0.00749238]; G_loss: [-0.00476407]\n",
      "Iter-79000; D_loss: [-0.00773986]; G_loss: [-0.00650107]\n",
      "Iter-80000; D_loss: [-0.00990202]; G_loss: [-0.00910741]\n",
      "Iter-81000; D_loss: [-0.00766142]; G_loss: [-0.00970357]\n",
      "Iter-82000; D_loss: [-0.00929849]; G_loss: [-0.00740123]\n",
      "Iter-83000; D_loss: [-0.00784644]; G_loss: [-0.00708846]\n",
      "Iter-84000; D_loss: [-0.00864486]; G_loss: [-0.00467691]\n",
      "Iter-85000; D_loss: [-0.01044832]; G_loss: [-0.00375227]\n",
      "Iter-86000; D_loss: [-0.01236129]; G_loss: [-0.00910825]\n",
      "Iter-87000; D_loss: [-0.00974829]; G_loss: [-0.00909376]\n",
      "Iter-88000; D_loss: [-0.0079511]; G_loss: [-0.00283779]\n",
      "Iter-89000; D_loss: [-0.00869218]; G_loss: [-0.00589232]\n",
      "Iter-90000; D_loss: [-0.00682809]; G_loss: [-0.00506519]\n",
      "Iter-91000; D_loss: [-0.00842675]; G_loss: [-0.00775276]\n",
      "Iter-92000; D_loss: [-0.00846413]; G_loss: [-0.00396187]\n",
      "Iter-93000; D_loss: [-0.00858897]; G_loss: [-0.00543921]\n",
      "Iter-94000; D_loss: [-0.00699276]; G_loss: [-0.00394842]\n",
      "Iter-95000; D_loss: [-0.00834101]; G_loss: [-0.00686648]\n",
      "Iter-96000; D_loss: [-0.00728609]; G_loss: [-0.00391402]\n",
      "Iter-97000; D_loss: [-0.00993953]; G_loss: [-0.00616756]\n",
      "Iter-98000; D_loss: [-0.01117188]; G_loss: [-0.00873939]\n",
      "Iter-99000; D_loss: [-0.00984632]; G_loss: [-0.00718678]\n",
      "Iter-100000; D_loss: [-0.0070178]; G_loss: [-0.00373726]\n",
      "Iter-101000; D_loss: [-0.00739946]; G_loss: [-0.00500223]\n",
      "Iter-102000; D_loss: [-0.00679002]; G_loss: [-0.00539502]\n",
      "Iter-103000; D_loss: [-0.00894472]; G_loss: [-0.00389481]\n",
      "Iter-104000; D_loss: [-0.00748858]; G_loss: [-0.00402257]\n",
      "Iter-105000; D_loss: [-0.00908742]; G_loss: [-0.00661718]\n",
      "Iter-106000; D_loss: [-0.00847542]; G_loss: [-0.00475244]\n",
      "Iter-107000; D_loss: [-0.00767949]; G_loss: [-0.00833201]\n",
      "Iter-108000; D_loss: [-0.00977358]; G_loss: [-0.00178068]\n",
      "Iter-109000; D_loss: [-0.008047]; G_loss: [-0.00074986]\n",
      "Iter-110000; D_loss: [-0.00918459]; G_loss: [-0.00161338]\n",
      "Iter-111000; D_loss: [-0.00840778]; G_loss: [-0.00444629]\n",
      "Iter-112000; D_loss: [-0.00563688]; G_loss: [-0.00825442]\n",
      "Iter-113000; D_loss: [-0.0081655]; G_loss: [-0.0056976]\n"
     ]
    }
   ],
   "source": [
    "for it in range(1000000):\n",
    "    for _ in range(5):\n",
    "        # Sample data\n",
    "        z = Variable(torch.randn(32, 10))\n",
    "        X, _ = mnist.train.next_batch(32)\n",
    "        X = Variable(torch.from_numpy(X))\n",
    "\n",
    "        # Dicriminator forward-loss-backward-update\n",
    "        G_sample = G(z)\n",
    "        D_real = D(X)\n",
    "        D_fake = D(G_sample)\n",
    "\n",
    "        D_loss = -(torch.mean(D_real) - torch.mean(D_fake))\n",
    "\n",
    "        D_loss.backward()\n",
    "        D_solver.step()\n",
    "\n",
    "        # Weight clipping\n",
    "        for p in D.parameters():\n",
    "            p.data.clamp_(-0.01, 0.01)\n",
    "\n",
    "        # Housekeeping - reset gradient\n",
    "        reset_grad()\n",
    "\n",
    "    # Generator forward-loss-backward-update\n",
    "    X, _ = mnist.train.next_batch(32)\n",
    "    X = Variable(torch.from_numpy(X))\n",
    "    z = Variable(torch.randn(32, 10))\n",
    "\n",
    "    G_sample = G(z)\n",
    "    D_fake = D(G_sample)\n",
    "\n",
    "    G_loss = -torch.mean(D_fake)\n",
    "\n",
    "    G_loss.backward()\n",
    "    G_solver.step()\n",
    "\n",
    "    # Housekeeping - reset gradient\n",
    "    reset_grad()\n",
    "\n",
    "    # Print and plot every now and then\n",
    "    if it % 1000 == 0:\n",
    "        print('Iter-{}; D_loss: {}; G_loss: {}'\n",
    "              .format(it, D_loss.data.numpy(), G_loss.data.numpy()))\n",
    "\n",
    "        samples = G(z).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'.format(str(cnt).zfill(3)), bbox_inches='tight')\n",
    "        cnt += 1\n",
    "        plt.close(fig)"
   ]
  },
  {
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}