WGAN_ADNI_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab inline\n",
    "import numpy as np\n",
    "import torch\n",
    "import os\n",
    "\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.nn import functional as F\n",
    "from torch import autograd\n",
    "from torch.autograd import Variable\n",
    "import nibabel as nib\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data import dataloader\n",
    "from nilearn import plotting\n",
    "from ADNI_dataset import *\n",
    "from BRATS_dataset import *\n",
    "from ATLAS_dataset import *\n",
    "from Model_WGAN import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE=4\n",
    "max_epoch = 100\n",
    "lr = 0.0001\n",
    "gpu = True\n",
    "workers = 4\n",
    "\n",
    "LAMBDA= 10\n",
    "#setting latent variable sizes\n",
    "latent_dim = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = ADNIdataset(augmentation=True)\n",
    "train_loader = torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True,num_workers=workers)\n",
    "if Use_BRATS:\n",
    "    #'flair' or 't2' or 't1ce'\n",
    "    trainset = BRATSdataset(imgtype='flair')\n",
    "    train_loader = torch.utils.data.DataLoader(trainset,batch_size = BATCH_SIZE, shuffle=True,\n",
    "                                               num_workers=workers)\n",
    "if Use_ATLAS:\n",
    "    trainset = ATLASdataset(augmentation=True)\n",
    "    train_loader = torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True,num_workers=workers)\n",
    "\n",
    "def inf_train_gen(data_loader):\n",
    "    while True:\n",
    "        for _,images in enumerate(data_loader):\n",
    "            yield images\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = Discriminator()\n",
    "G = Generator(noise = latent_dim)\n",
    "\n",
    "g_optimizer = optim.Adam(G.parameters(), lr=0.0002)\n",
    "d_optimizer = optim.Adam(D.parameters(), lr=0.0002)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_gradient_penalty(netD, real_data, fake_data):    \n",
    "    alpha = torch.rand(real_data.size(0),1,1,1,1)\n",
    "    alpha = alpha.expand(real_data.size())\n",
    "    \n",
    "    alpha = alpha.cuda()\n",
    "\n",
    "    interpolates = alpha * real_data + ((1 - alpha) * fake_data)\n",
    "\n",
    "    interpolates = interpolates.cuda()\n",
    "    interpolates = Variable(interpolates, requires_grad=True)\n",
    "\n",
    "    disc_interpolates = netD(interpolates)\n",
    "\n",
    "    gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,\n",
    "                              grad_outputs=torch.ones(disc_interpolates.size()).cuda(),\n",
    "                              create_graph=True, retain_graph=True, only_inputs=True)[0]\n",
    "\n",
    "    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * LAMBDA\n",
    "    return gradient_penalty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_y = Variable(torch.ones((BATCH_SIZE, 1)).cuda())\n",
    "fake_y = Variable(torch.zeros((BATCH_SIZE, 1)).cuda())\n",
    "loss_f = nn.BCELoss()\n",
    "\n",
    "d_real_losses = list()\n",
    "d_fake_losses = list()\n",
    "d_losses = list()\n",
    "g_losses = list()\n",
    "divergences = list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "TOTAL_ITER = 200000\n",
    "gen_load = inf_train_gen(train_loader)\n",
    "for iteration in range(TOTAL_ITER):\n",
    "    ###############################################\n",
    "    # Train D \n",
    "    ###############################################\n",
    "    for p in D.parameters():  \n",
    "        p.requires_grad = True \n",
    "\n",
    "    real_images = gen_load.__next__()\n",
    "    D.zero_grad()\n",
    "    real_images = Variable(real_images).cuda()\n",
    "\n",
    "    _batch_size = real_images.size(0)\n",
    "\n",
    "\n",
    "    y_real_pred = D(real_images)\n",
    "\n",
    "    d_real_loss = y_real_pred.mean()\n",
    "    \n",
    "    noise = Variable(torch.randn((_batch_size, latent_dim, 1, 1, 1)),volatile=True).cuda()\n",
    "    fake_images = G(noise)\n",
    "    y_fake_pred = D(fake_images.detach())\n",
    "\n",
    "    d_fake_loss = y_fake_pred.mean()\n",
    "\n",
    "    gradient_penalty = calc_gradient_penalty(D,real_images.data, fake_images.data)\n",
    " \n",
    "    d_loss = - d_real_loss + d_fake_loss +gradient_penalty\n",
    "    d_loss.backward()\n",
    "    Wasserstein_D = d_real_loss - d_fake_loss\n",
    "\n",
    "    d_optimizer.step()\n",
    "\n",
    "    ###############################################\n",
    "    # Train G \n",
    "    ###############################################\n",
    "    for p in D.parameters():\n",
    "        p.requires_grad = False\n",
    "        \n",
    "    for iters in range(5):\n",
    "        G.zero_grad()\n",
    "        noise = Variable(torch.randn((_batch_size, latent_dim, 1, 1 ,1)).cuda())\n",
    "        fake_image =G(noise)\n",
    "        y_fake_g = D(fake_image)\n",
    "\n",
    "        g_loss = -y_fake_g.mean()\n",
    "\n",
    "        g_loss.backward()\n",
    "        g_optimizer.step()\n",
    "\n",
    "    ###############################################\n",
    "    # Visualization\n",
    "    ###############################################\n",
    "    if iteration%10 == 0:\n",
    "        d_real_losses.append(d_real_loss.data[0])\n",
    "        d_fake_losses.append(d_fake_loss.data[0])\n",
    "        d_losses.append(d_loss.data[0])\n",
    "        g_losses.append(g_loss.data.cpu().numpy())\n",
    "\n",
    "        print('[{}/{}]'.format(iteration,TOTAL_ITER),\n",
    "              'D: {:<8.3}'.format(d_loss.data[0].cpu().numpy()), \n",
    "              'D_real: {:<8.3}'.format(d_real_loss.data[0].cpu().numpy()),\n",
    "              'D_fake: {:<8.3}'.format(d_fake_loss.data[0].cpu().numpy()), \n",
    "              'G: {:<8.3}'.format(g_loss.data[0].cpu().numpy()))\n",
    "\n",
    "        featmask = np.squeeze((0.5*fake_image+0.5)[0].data.cpu().numpy())\n",
    "\n",
    "        featmask = nib.Nifti1Image(featmask,affine = np.eye(4))\n",
    "        plotting.plot_img(featmask,title=\"FAKE\")\n",
    "        plotting.show()\n",
    "        \n",
    "    if (iteration+1)%500 ==0:\n",
    "        torch.save(G.state_dict(),'./checkpoint/G_W_iter'+str(iteration+1)+'.pth')\n",
    "        torch.save(D.state_dict(),'./checkpoint/D_W_iter'+str(iteration+1)+'.pth')"
   ]
  }
 ],
 "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.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}