2. DC-GAN PyTorch-MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython import display\n",
    "\n",
    "from utils import Logger\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.optim import Adam\n",
    "from torch.autograd import Variable\n",
    "\n",
    "from torchvision import transforms, datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_FOLDER = './torch_data/DCGAN/MNIST'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_data():\n",
    "    compose = transforms.Compose(\n",
    "        [\n",
    "            transforms.Resize(64),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((.5,), (.5,))\n",
    "        ])\n",
    "    out_dir = '{}/dataset'.format(DATA_FOLDER)\n",
    "    return datasets.MNIST(root=out_dir, train=True, transform=compose, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = mnist_data()\n",
    "batch_size = 100\n",
    "data_loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=True)\n",
    "num_batches = len(data_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiscriminativeNet(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(DiscriminativeNet, self).__init__()\n",
    "        \n",
    "        self.conv1 = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=1, out_channels=128, kernel_size=4, \n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        self.conv2 = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=128, out_channels=256, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        self.conv3 = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=256, out_channels=512, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        self.conv4 = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=512, out_channels=1024, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        self.out = nn.Sequential(\n",
    "            nn.Linear(1024*4*4, 1),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Convolutional layers\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.conv4(x)\n",
    "        # Flatten and apply sigmoid\n",
    "        x = x.view(-1, 1024*4*4)\n",
    "        x = self.out(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GenerativeNet(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(GenerativeNet, self).__init__()\n",
    "        \n",
    "        self.linear = torch.nn.Linear(100, 1024*4*4)\n",
    "        \n",
    "        self.conv1 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=1024, out_channels=512, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        self.conv2 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=512, out_channels=256, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        self.conv3 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=256, out_channels=128, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            ),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        self.conv4 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=128, out_channels=1, kernel_size=4,\n",
    "                stride=2, padding=1, bias=False\n",
    "            )\n",
    "        )\n",
    "        self.out = torch.nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Project and reshape\n",
    "        x = self.linear(x)\n",
    "        x = x.view(x.shape[0], 1024, 4, 4)\n",
    "        # Convolutional layers\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.conv4(x)\n",
    "        # Apply Tanh\n",
    "        return self.out(x)\n",
    "    \n",
    "# Noise\n",
    "def noise(size):\n",
    "    n = Variable(torch.randn(size, 100))\n",
    "    if torch.cuda.is_available(): return n.cuda()\n",
    "    return n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1 or classname.find('BatchNorm') != -1:\n",
    "        m.weight.data.normal_(0.00, 0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Network instances and init weights\n",
    "generator = GenerativeNet()\n",
    "generator.apply(init_weights)\n",
    "\n",
    "discriminator = DiscriminativeNet()\n",
    "discriminator.apply(init_weights)\n",
    "\n",
    "# Enable cuda if available\n",
    "if torch.cuda.is_available():\n",
    "    generator.cuda()\n",
    "    discriminator.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizers\n",
    "d_optimizer = Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "g_optimizer = Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "\n",
    "# Loss function\n",
    "loss = nn.BCELoss()\n",
    "\n",
    "# Number of epochs\n",
    "num_epochs = 200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def real_data_target(size):\n",
    "    '''\n",
    "    Tensor containing ones, with shape = size\n",
    "    '''\n",
    "    data = Variable(torch.ones(size, 1))\n",
    "    if torch.cuda.is_available(): return data.cuda()\n",
    "    return data\n",
    "\n",
    "def fake_data_target(size):\n",
    "    '''\n",
    "    Tensor containing zeros, with shape = size\n",
    "    '''\n",
    "    data = Variable(torch.zeros(size, 1))\n",
    "    if torch.cuda.is_available(): return data.cuda()\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_discriminator(optimizer, real_data, fake_data):\n",
    "    # Reset gradients\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # 1. Train on Real Data\n",
    "    prediction_real = discriminator(real_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error_real = loss(prediction_real, real_data_target(real_data.size(0)))\n",
    "    error_real.backward()\n",
    "\n",
    "    # 2. Train on Fake Data\n",
    "    prediction_fake = discriminator(fake_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error_fake = loss(prediction_fake, fake_data_target(real_data.size(0)))\n",
    "    error_fake.backward()\n",
    "    \n",
    "    # Update weights with gradients\n",
    "    optimizer.step()\n",
    "    \n",
    "    return error_real + error_fake, prediction_real, prediction_fake\n",
    "    return (0, 0, 0)\n",
    "\n",
    "def train_generator(optimizer, fake_data):\n",
    "    # Reset gradients\n",
    "    optimizer.zero_grad()\n",
    "    # Sample noise and generate fake data\n",
    "    prediction = discriminator(fake_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error = loss(prediction, real_data_target(prediction.size(0)))\n",
    "    error.backward()\n",
    "    # Update weights with gradients\n",
    "    optimizer.step()\n",
    "    # Return error\n",
    "    return error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Samples for Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_test_samples = 16\n",
    "test_noise = noise(num_test_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "logger = Logger(model_name='DCGAN', data_name='MNIST')\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    for n_batch, (real_batch,_) in enumerate(data_loader):\n",
    "        \n",
    "        # 1. Train Discriminator\n",
    "        real_data = Variable(real_batch)\n",
    "        if torch.cuda.is_available(): real_data = real_data.cuda()\n",
    "        # Generate fake data\n",
    "        fake_data = generator(noise(real_data.size(0))).detach()\n",
    "        # Train D\n",
    "        d_error, d_pred_real, d_pred_fake = train_discriminator(d_optimizer, \n",
    "                                                                real_data, fake_data)\n",
    "\n",
    "        # 2. Train Generator\n",
    "        # Generate fake data\n",
    "        fake_data = generator(noise(real_batch.size(0)))\n",
    "        # Train G\n",
    "        g_error = train_generator(g_optimizer, fake_data)\n",
    "        # Log error\n",
    "        logger.log(d_error, g_error, epoch, n_batch, num_batches)\n",
    "        \n",
    "        # Display Progress\n",
    "        if (n_batch) % 100 == 0:\n",
    "            display.clear_output(True)\n",
    "            # Display Images\n",
    "            test_images = generator(test_noise).data.cpu()\n",
    "            logger.log_images(test_images, num_test_samples, epoch, n_batch, num_batches);\n",
    "            # Display status Logs\n",
    "            logger.display_status(\n",
    "                epoch, num_epochs, n_batch, num_batches,\n",
    "                d_error, g_error, d_pred_real, d_pred_fake\n",
    "            )\n",
    "        # Model Checkpoints\n",
    "        logger.save_models(generator, discriminator, epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}