Pytorch_DCGAN_CIFAR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import os\n",
    "import random\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.utils as vutils\n",
    "\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Init args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self):\n",
    "        self.dataset=\"cifar10\" #cifar10 | lsun | mnist |imagenet | folder | lfw | fake\n",
    "        self.dataroot=\".\"\n",
    "        self.workers=2 #number of data loading workers\n",
    "        self.batchSize=64\n",
    "        self.imageSize=64\n",
    "        self.cuda=True\n",
    "        self.nz=100 #size of the latent z vector\n",
    "        self.ngf=64\n",
    "        self.ndf=64\n",
    "        self.epochs=3 #number of epochs to train for\n",
    "        self.lr=0.0002\n",
    "        self.beta1=0.5 #beta1 for adam\n",
    "        self.ngpu=1 #number of GPUs to use\n",
    "        self.netG=\"\" #path to netG (to continue training)\n",
    "        self.netD=\"\" #path to netD (to continue training)\n",
    "        self.outf=\".\" #folder to output images and model checkpoints\n",
    "        self.manualSeed=1 #manual seed\n",
    "        self.classes=\"bedroom\" #comma separated list of classes for the lsun data set\n",
    "\n",
    "opt=Args()\n",
    "try:\n",
    "    os.makedirs(opt.outf)\n",
    "except OSError:\n",
    "    pass\n",
    "if opt.manualSeed is None:\n",
    "    opt.manualSeed = random.randint(1, 10000)\n",
    "print(\"Random Seed: \", opt.manualSeed)\n",
    "random.seed(opt.manualSeed)\n",
    "torch.manual_seed(opt.manualSeed)\n",
    "\n",
    "device = torch.device(\"cuda:0\" if opt.cuda else \"cpu\")\n",
    "ngpu = int(opt.ngpu)\n",
    "nz = int(opt.nz)\n",
    "ngf = int(opt.ngf)\n",
    "ndf = int(opt.ndf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if opt.dataset in ['imagenet', 'folder', 'lfw']:\n",
    "    # folder dataset\n",
    "    dataset = dset.ImageFolder(root=opt.dataroot,\n",
    "                               transform=transforms.Compose([\n",
    "                                   transforms.Resize(opt.imageSize),\n",
    "                                   transforms.CenterCrop(opt.imageSize),\n",
    "                                   transforms.ToTensor(),\n",
    "                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "                               ]))\n",
    "    nc=3\n",
    "elif opt.dataset == 'lsun':\n",
    "    classes = [ c + '_train' for c in opt.classes.split(',')]\n",
    "    dataset = dset.LSUN(root=opt.dataroot, classes=classes,\n",
    "                        transform=transforms.Compose([\n",
    "                            transforms.Resize(opt.imageSize),\n",
    "                            transforms.CenterCrop(opt.imageSize),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "                        ]))\n",
    "    nc=3\n",
    "elif opt.dataset == 'cifar10':\n",
    "    dataset = dset.CIFAR10(root=opt.dataroot, download=True,\n",
    "                           transform=transforms.Compose([\n",
    "                               transforms.Resize(opt.imageSize),\n",
    "                               transforms.ToTensor(),\n",
    "                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "                           ]))\n",
    "    nc=3\n",
    "\n",
    "elif opt.dataset == 'mnist':\n",
    "    dataset = dset.MNIST(root=opt.dataroot, download=True,\n",
    "                       transform=transforms.Compose([\n",
    "                           transforms.Resize(opt.imageSize),\n",
    "                           transforms.ToTensor(),\n",
    "                           transforms.Normalize((0.5,), (0.5,)),\n",
    "                       ]))\n",
    "    nc=1\n",
    "\n",
    "elif opt.dataset == 'fake':\n",
    "    dataset = dset.FakeData(image_size=(3, opt.imageSize, opt.imageSize),\n",
    "                            transform=transforms.ToTensor())\n",
    "    nc=3\n",
    "\n",
    "assert dataset\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,\n",
    "                                         shuffle=True, num_workers=int(opt.workers))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Init model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# custom weights initialization called on netG and netD\n",
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        m.weight.data.normal_(0.0, 0.02)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        m.weight.data.normal_(1.0, 0.02)\n",
    "        m.bias.data.fill_(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Generator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.main = nn.Sequential(\n",
    "            # input is Z, going into a convolution\n",
    "            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 8),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*8) x 4 x 4\n",
    "            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 4),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*4) x 8 x 8\n",
    "            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf*2) x 16 x 16\n",
    "            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (ngf) x 32 x 32\n",
    "            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),\n",
    "            nn.Tanh()\n",
    "            # state size. (nc) x 64 x 64\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        if input.is_cuda and self.ngpu > 1:\n",
    "            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))\n",
    "        else:\n",
    "            output = self.main(input)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.main = nn.Sequential(\n",
    "            # input is (nc) x 64 x 64\n",
    "            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # state size. (ndf) x 32 x 32\n",
    "            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 2),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # state size. (ndf*2) x 16 x 16\n",
    "            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # state size. (ndf*4) x 8 x 8\n",
    "            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 8),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # state size. (ndf*8) x 4 x 4\n",
    "            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        if input.is_cuda and self.ngpu > 1:\n",
    "            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))\n",
    "        else:\n",
    "            output = self.main(input)\n",
    "\n",
    "        return output.view(-1, 1).squeeze(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "netG = Generator(ngpu).to(device)\n",
    "netG.apply(weights_init)\n",
    "if opt.netG != '':\n",
    "    netG.load_state_dict(torch.load(opt.netG))\n",
    "print(netG)\n",
    "\n",
    "netD = Discriminator(ngpu).to(device)\n",
    "netD.apply(weights_init)\n",
    "if opt.netD != '':\n",
    "    netD.load_state_dict(torch.load(opt.netD))\n",
    "print(netD)\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "\n",
    "fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device)\n",
    "real_label = 1\n",
    "fake_label = 0\n",
    "\n",
    "# setup optimizer\n",
    "optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))\n",
    "optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(opt.epochs):\n",
    "    for i, data in enumerate(dataloader, 0):\n",
    "        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "        # train with real\n",
    "        netD.zero_grad()\n",
    "        real_cpu = data[0].to(device)\n",
    "        batch_size = real_cpu.size(0)\n",
    "        label = torch.full((batch_size,), real_label, device=device)\n",
    "\n",
    "        output = netD(real_cpu)\n",
    "        errD_real = criterion(output, label)\n",
    "        errD_real.backward()\n",
    "        D_x = output.mean().item()\n",
    "\n",
    "        # train with fake\n",
    "        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n",
    "        fake = netG(noise)\n",
    "        label.fill_(fake_label)\n",
    "        output = netD(fake.detach())\n",
    "        errD_fake = criterion(output, label)\n",
    "        errD_fake.backward()\n",
    "        D_G_z1 = output.mean().item()\n",
    "        errD = errD_real + errD_fake\n",
    "        optimizerD.step()\n",
    "\n",
    "        # (2) Update G network: maximize log(D(G(z)))\n",
    "        netG.zero_grad()\n",
    "        label.fill_(real_label)  # fake labels are real for generator cost\n",
    "        output = netD(fake)\n",
    "        errG = criterion(output, label)\n",
    "        errG.backward()\n",
    "        D_G_z2 = output.mean().item()\n",
    "        optimizerG.step()\n",
    "\n",
    "    print('[%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'\n",
    "          % (epoch, opt.epochs,errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
    "    vutils.save_image(real_cpu,\n",
    "            '%s/real_samples.png' % opt.outf,\n",
    "            normalize=True)\n",
    "    fake = netG(fixed_noise)\n",
    "    vutils.save_image(fake.detach(),\n",
    "            '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),\n",
    "            normalize=True)\n",
    "    # do checkpointing\n",
    "    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))\n",
    "    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}