Alpha_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",
    "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_alphaWGAN import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE=4\n",
    "gpu = True\n",
    "workers = 4\n",
    "\n",
    "LAMBDA= 10\n",
    "_eps = 1e-15\n",
    "Use_BRATS=False\n",
    "Use_ATLAS = False\n",
    "\n",
    "#setting latent variable sizes\n",
    "latent_dim = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inf_train_gen(data_loader):\n",
    "    while True:\n",
    "        for _,images in enumerate(data_loader):\n",
    "            yield images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = Generator(noise = latent_dim)\n",
    "CD = Code_Discriminator(code_size = latent_dim ,num_units = 4096)\n",
    "D = Discriminator(is_dis=True)\n",
    "E = Discriminator(out_class = latent_dim,is_dis=False)\n",
    "\n",
    "G.cuda()\n",
    "D.cuda()\n",
    "CD.cuda()\n",
    "E.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_optimizer = optim.Adam(G.parameters(), lr=0.0002)\n",
    "d_optimizer = optim.Adam(D.parameters(), lr=0.0002)\n",
    "e_optimizer = optim.Adam(E.parameters(), lr = 0.0002)\n",
    "cd_optimizer = optim.Adam(CD.parameters(), lr = 0.0002)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_gradient_penalty(model, x, x_gen, w=10):\n",
    "    \"\"\"WGAN-GP gradient penalty\"\"\"\n",
    "    assert x.size()==x_gen.size(), \"real and sampled sizes do not match\"\n",
    "    alpha_size = tuple((len(x), *(1,)*(x.dim()-1)))\n",
    "    alpha_t = torch.cuda.FloatTensor if x.is_cuda else torch.Tensor\n",
    "    alpha = alpha_t(*alpha_size).uniform_()\n",
    "    x_hat = x.data*alpha + x_gen.data*(1-alpha)\n",
    "    x_hat = Variable(x_hat, requires_grad=True)\n",
    "\n",
    "    def eps_norm(x):\n",
    "        x = x.view(len(x), -1)\n",
    "        return (x*x+_eps).sum(-1).sqrt()\n",
    "    def bi_penalty(x):\n",
    "        return (x-1)**2\n",
    "\n",
    "    grad_xhat = torch.autograd.grad(model(x_hat).sum(), x_hat, create_graph=True, only_inputs=True)[0]\n",
    "\n",
    "    penalty = w*bi_penalty(eps_norm(grad_xhat)).mean()\n",
    "    return 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(async=True))\n",
    "fake_y = Variable(torch.zeros((BATCH_SIZE, 1)).cuda(async=True))\n",
    "\n",
    "criterion_bce = nn.BCELoss()\n",
    "criterion_l1 = nn.L1Loss()\n",
    "criterion_mse = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_iter = 1\n",
    "d_iter = 1\n",
    "cd_iter =1\n",
    "TOTAL_ITER = 200000\n",
    "gen_load = inf_train_gen(train_loader)\n",
    "for iteration in range(TOTAL_ITER):\n",
    "    for p in D.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in CD.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in E.parameters():  \n",
    "        p.requires_grad = True\n",
    "    for p in G.parameters():  \n",
    "        p.requires_grad = True\n",
    "\n",
    "    ###############################################\n",
    "    # Train Encoder - Generator \n",
    "    ###############################################\n",
    "    for iters in range(g_iter):\n",
    "        G.zero_grad()\n",
    "        E.zero_grad()\n",
    "        real_images = gen_load.__next__()\n",
    "        _batch_size = real_images.size(0)\n",
    "        real_images = Variable(real_images,volatile=True).cuda(async=True)\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),volatile=True).cuda()\n",
    "        z_hat = E(real_images).view(_batch_size,-1)\n",
    "        x_hat = G(z_hat)\n",
    "        x_rand = G(z_rand)\n",
    "        c_loss = -CD(z_hat).mean()\n",
    "\n",
    "        d_real_loss = D(x_hat).mean()\n",
    "        d_fake_loss = D(x_rand).mean()\n",
    "        d_loss = -d_fake_loss-d_real_loss\n",
    "        l1_loss =10* criterion_l1(x_hat,real_images)\n",
    "        loss1 = l1_loss + c_loss + d_loss\n",
    "\n",
    "        if iters<g_iter-1:\n",
    "            loss1.backward()\n",
    "        else:\n",
    "            loss1.backward(retain_graph=True)\n",
    "        e_optimizer.step()\n",
    "        g_optimizer.step()\n",
    "        g_optimizer.step()\n",
    "\n",
    "\n",
    "    ###############################################\n",
    "    # Train D\n",
    "    ###############################################\n",
    "    for p in D.parameters():  \n",
    "        p.requires_grad = True\n",
    "    for p in CD.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in E.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in G.parameters():  \n",
    "        p.requires_grad = False\n",
    "\n",
    "    for iters in range(d_iter):\n",
    "        d_optimizer.zero_grad()\n",
    "        real_images = gen_load.__next__()\n",
    "        _batch_size = real_images.size(0)\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),volatile=True).cuda()\n",
    "        real_images = Variable(real_images,volatile=True).cuda(async=True)\n",
    "        z_hat = E(real_images).view(_batch_size,-1)\n",
    "        x_hat = G(z_hat)\n",
    "        x_rand = G(z_rand)\n",
    "        x_loss2 = -2*D(real_images).mean()+D(x_hat).mean()+D(x_rand).mean()\n",
    "        gradient_penalty_r = calc_gradient_penalty(D,real_images.data, x_rand.data)\n",
    "        gradient_penalty_h = calc_gradient_penalty(D,real_images.data, x_hat.data)\n",
    "\n",
    "        loss2 = x_loss2+gradient_penalty_r+gradient_penalty_h\n",
    "        loss2.backward(retain_graph=True)\n",
    "        d_optimizer.step()\n",
    "\n",
    "    ###############################################\n",
    "    # Train CD\n",
    "    ###############################################\n",
    "    for p in D.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in CD.parameters():  \n",
    "        p.requires_grad = True\n",
    "    for p in E.parameters():  \n",
    "        p.requires_grad = False\n",
    "    for p in G.parameters():  \n",
    "        p.requires_grad = False\n",
    "\n",
    "    for iters in range(cd_iter):\n",
    "        cd_optimizer.zero_grad()\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),volatile=True).cuda()\n",
    "        gradient_penalty_cd = calc_gradient_penalty(CD,z_hat.data, z_rand.data)\n",
    "        loss3 = -CD(z_rand).mean() - c_loss + gradient_penalty_cd\n",
    "\n",
    "        loss3.backward(retain_graph=True)\n",
    "        cd_optimizer.step()\n",
    "\n",
    "    ###############################################\n",
    "    # Visualization\n",
    "    ###############################################\n",
    "\n",
    "    if iteration % 10 == 0:\n",
    "        print('[{}/{}]'.format(iteration,TOTAL_ITER),\n",
    "              'D: {:<8.3}'.format(loss2.data[0].cpu().numpy()), \n",
    "              'En_Ge: {:<8.3}'.format(loss1.data[0].cpu().numpy()),\n",
    "              'Code: {:<8.3}'.format(loss3.data[0].cpu().numpy()),\n",
    "              )\n",
    "        feat = np.squeeze((0.5*real_images[0]+0.5).data.cpu().numpy())\n",
    "        feat = nib.Nifti1Image(feat,affine = np.eye(4))\n",
    "        plotting.plot_img(feat,title=\"X_Real\")\n",
    "        plotting.show()\n",
    "\n",
    "        feat = np.squeeze((0.5*x_hat[0]+0.5).data.cpu().numpy())\n",
    "        feat = nib.Nifti1Image(feat,affine = np.eye(4))\n",
    "        plotting.plot_img(feat,title=\"X_DEC\")\n",
    "        plotting.show()\n",
    "\n",
    "        feat = np.squeeze((0.5*x_rand[0]+0.5).data.cpu().numpy())\n",
    "        feat = nib.Nifti1Image(feat,affine = np.eye(4))\n",
    "        plotting.plot_img(feat,title=\"X_rand\")\n",
    "        plotting.show()\n",
    "\n",
    "    ###############################################\n",
    "    # Model Save\n",
    "    ###############################################\n",
    "    if (iteration+1)%500 ==0: \n",
    "        torch.save(G.state_dict(),'./checkpoint/G_iter'+str(iteration+1+es)+'.pth')\n",
    "        torch.save(D.state_dict(),'./checkpoint/D_iter'+str(iteration+1+es)+'.pth')\n",
    "        torch.save(E.state_dict(),'./checkpoint/E_iter'+str(iteration+1+es)+'.pth')\n",
    "        torch.save(CD.state_dict(),'./checkpoint/CD_iter'+str(iteration+1+es)+'.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
}