VAEGAN_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 skimage.transform import resize\n",
    "from nilearn import plotting\n",
    "from ADNI_dataset import *\n",
    "from BRATS_dataset import *\n",
    "from ATLAS_dataset import *\n",
    "from Model_VAEGAN import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE=4\n",
    "max_epoch = 100\n",
    "gpu = True\n",
    "workers = 4\n",
    "\n",
    "reg = 5e-10\n",
    "\n",
    "gamma = 20\n",
    "beta = 10\n",
    "\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": {},
   "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": [
    "G = Generator(noise = latent_dim)\n",
    "D = Discriminator()\n",
    "E = Encoder()\n",
    "\n",
    "G.cuda()\n",
    "D.cuda()\n",
    "E.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_optimizer = optim.Adam(G.parameters(), lr=0.0001)\n",
    "d_optimizer = optim.Adam(D.parameters(), lr=0.0001)\n",
    "e_optimizer = optim.Adam(E.parameters(), lr = 0.0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCH = 100\n",
    "\n",
    "real_y = Variable(torch.ones((BATCH_SIZE, 1)).cuda())\n",
    "fake_y = Variable(torch.zeros((BATCH_SIZE, 1)).cuda())\n",
    "criterion_bce = nn.BCELoss()\n",
    "criterion_l1 = nn.L1Loss()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for epoch in range(N_EPOCH):\n",
    "    for step, real_images in enumerate(train_loader):\n",
    "        _batch_size = real_images.size(0)\n",
    "        real_images = Variable(real_images,requires_grad=False).cuda()\n",
    "        z_rand = Variable(torch.randn((_batch_size, latent_dim)),requires_grad=False).cuda()\n",
    "        mean,logvar,code = E(real_images)\n",
    "        x_rec = G(code)\n",
    "        x_rand = G(z_rand)\n",
    "        ###############################################\n",
    "        # Train D \n",
    "        ###############################################\n",
    "        d_optimizer.zero_grad()\n",
    "        \n",
    "        d_real_loss = criterion_bce(D(real_images),real_y[:_batch_size])\n",
    "        d_recon_loss = criterion_bce(D(x_rec), fake_y[:_batch_size])\n",
    "        d_fake_loss = criterion_bce(D(x_rand), fake_y[:_batch_size])\n",
    "        \n",
    "        dis_loss = d_recon_loss+d_real_loss + d_fake_loss\n",
    "        dis_loss.backward(retain_graph=True)\n",
    "        \n",
    "        d_optimizer.step()\n",
    "        \n",
    "        ###############################################\n",
    "        # Train G\n",
    "        ###############################################\n",
    "        g_optimizer.zero_grad()\n",
    "        output = D(real_images)\n",
    "        d_real_loss = criterion_bce(output,real_y[:_batch_size])\n",
    "        output = D(x_rec)\n",
    "        d_recon_loss = criterion_bce(output,fake_y[:_batch_size])\n",
    "        output = D(x_rand)\n",
    "        d_fake_loss = criterion_bce(output,fake_y[:_batch_size])\n",
    "        \n",
    "        d_img_loss = d_real_loss + d_recon_loss+ d_fake_loss\n",
    "        gen_img_loss = -d_img_loss\n",
    "        \n",
    "        rec_loss = ((x_rec - real_images)**2).mean()\n",
    "        \n",
    "        err_dec = gamma* rec_loss + gen_img_loss\n",
    "        \n",
    "        err_dec.backward(retain_graph=True)\n",
    "        g_optimizer.step()\n",
    "        ###############################################\n",
    "        # Train E\n",
    "        ###############################################\n",
    "        prior_loss = 1+logvar-mean.pow(2) - logvar.exp()\n",
    "        prior_loss = (-0.5*torch.sum(prior_loss))/torch.numel(mean.data)\n",
    "        err_enc = prior_loss + beta*rec_loss\n",
    "        \n",
    "        e_optimizer.zero_grad()\n",
    "        err_enc.backward()\n",
    "        e_optimizer.step()\n",
    "        ###############################################\n",
    "        # Visualization\n",
    "        ###############################################\n",
    "  \n",
    "        if step % 10 == 0:\n",
    "            print('[{}/{}]'.format(epoch,N_EPOCH),\n",
    "                  'D: {:<8.3}'.format(dis_loss.data[0].cpu().numpy()), \n",
    "                  'En: {:<8.3}'.format(err_enc.data[0].cpu().numpy()),\n",
    "                  'De: {:<8.3}'.format(err_dec.data[0].cpu().numpy()) \n",
    "                  )\n",
    "            \n",
    "            featmask = np.squeeze((0.5*real_images[0]+0.5).data.cpu().numpy())\n",
    "            featmask = nib.Nifti1Image(featmask,affine = np.eye(4))\n",
    "            plotting.plot_img(featmask,title=\"X_Real\")\n",
    "            plotting.show()\n",
    "            \n",
    "            featmask = np.squeeze((0.5*x_rec[0]+0.5).data.cpu().numpy())\n",
    "            featmask = nib.Nifti1Image(featmask,affine = np.eye(4))\n",
    "            plotting.plot_img(featmask,title=\"X_DEC\")\n",
    "            plotting.show()\n",
    "            \n",
    "            featmask = np.squeeze((0.5*x_rand[0]+0.5).data.cpu().numpy())\n",
    "            featmask = nib.Nifti1Image(featmask,affine = np.eye(4))\n",
    "            plotting.plot_img(featmask,title=\"X_rand\")\n",
    "            plotting.show()\n",
    "\n",
    "    torch.save(G.state_dict(),'./chechpoint/G_VG_ep_'+str(epoch+1)+'.pth')\n",
    "    torch.save(D.state_dict(),'./chechpoint/D_VG_ep_'+str(epoch+1)+'.pth')\n",
    "    torch.save(E.state_dict(),'./chechpoint/E_VG_ep_'+str(epoch+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
}