Alpha_GAN_ADNI_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "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",
    "from ipdb import set_trace\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_alphaGAN import *\n",
    "from utils import *\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE=4\n",
    "gpu = True\n",
    "workers = 4\n",
    "LAMBDA= 10\n",
    "_eps = 1e-15\n",
    "device = torch.device('cuda:1')\n",
    "\n",
    "Use_BRATS = False\n",
    "Use_ATLAS = False\n",
    "\n",
    "#setting latent variable sizes\n",
    "latent_dim = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainset = ADNIdataset(augmentation=False)\n",
    "train_loader = torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True,num_workers=workers)\n",
    "\n",
    "if Use_BRATS:\n",
    "    #imgtype -> 'flair' or 't2' or 't1ce'\n",
    "    trainset = BRATSdataset(train=True, imgtype = 'flair',augmentation=False)\n",
    "    train_loader = torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True,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": 4,
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Discriminator(\n",
       "  (conv1): Conv3d(1, 64, kernel_size=(4, 4, 4), stride=(2, 2, 2), padding=(1, 1, 1))\n",
       "  (conv2): Conv3d(64, 128, kernel_size=(4, 4, 4), stride=(2, 2, 2), padding=(1, 1, 1))\n",
       "  (bn2): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3): Conv3d(128, 256, kernel_size=(4, 4, 4), stride=(2, 2, 2), padding=(1, 1, 1))\n",
       "  (bn3): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4): Conv3d(256, 512, kernel_size=(4, 4, 4), stride=(2, 2, 2), padding=(1, 1, 1))\n",
       "  (bn4): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5): Conv3d(512, 1000, kernel_size=(4, 4, 4), stride=(1, 1, 1))\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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.to(device)\n",
    "D.to(device)\n",
    "CD.to(device)\n",
    "E.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "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": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: './checkpoint/D_noW_iter0.pth'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-43057668a0f5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;31m# load the highest savepoints of all models\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m \u001b[0miteration\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_checkpoint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mG\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mE\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mCD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'_noW_iter'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0mdf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;31m# iteration = 0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/projects/shibo/3dbraingen/utils/__init__.py\u001b[0m in \u001b[0;36mload_checkpoint\u001b[0;34m(G, D, E, CD, fname)\u001b[0m\n\u001b[1;32m     16\u001b[0m                 \u001b[0mhighest_pth\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhighest_pth\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcurr_num\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mfiles\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m             \u001b[0mD\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'./checkpoint/D{fname}{highest_pth}.pth'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m             \u001b[0mCD\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'./checkpoint/CD{fname}{highest_pth}.pth'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m             \u001b[0mE\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'./checkpoint/E{fname}{highest_pth}.pth'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/miniconda3/lib/python3.7/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module)\u001b[0m\n\u001b[1;32m    363\u001b[0m             \u001b[0;34m(\u001b[0m\u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m3\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpathlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    364\u001b[0m         \u001b[0mnew_fd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 365\u001b[0;31m         \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    366\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    367\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_load\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmap_location\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpickle_module\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: './checkpoint/D_noW_iter0.pth'"
     ]
    }
   ],
   "source": [
    "real_y = Variable(torch.ones((BATCH_SIZE, 1)).to(device=device))\n",
    "fake_y = Variable(torch.zeros((BATCH_SIZE, 1)).to(device=device))\n",
    "\n",
    "criterion_bce = nn.BCELoss()\n",
    "criterion_l1 = nn.L1Loss()\n",
    "criterion_mse = nn.MSELoss()\n",
    "\n",
    "# load the highest savepoints of all models\n",
    "iteration = load_checkpoint(G, D, E, CD, '_noW_iter')\n",
    "df = load_loss()\n",
    "# iteration = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gen_load = inf_train_gen(train_loader)\n",
    "MAX_ITER = 200000\n",
    "while iteration < MAX_ITER:\n",
    "    ###############################################\n",
    "    # Train Encoder - Generator \n",
    "    ###############################################\n",
    "    for p in D.parameters():  # reset requires_grad\n",
    "        p.requires_grad = False\n",
    "    for p in CD.parameters():  # reset requires_grad\n",
    "        p.requires_grad = False\n",
    "    for p in E.parameters():  # reset requires_grad\n",
    "        p.requires_grad = True\n",
    "    for p in G.parameters():  # reset requires_grad\n",
    "        p.requires_grad = True\n",
    "\n",
    "    g_optimizer.zero_grad()\n",
    "    e_optimizer.zero_grad()\n",
    "\n",
    "\n",
    "    for iters in range(1):\n",
    "        real_images = gen_load.__next__()\n",
    "        real_images = Variable(real_images,volatile=True).to(device=device)\n",
    "        _batch_size = real_images.size(0)\n",
    "        z_hat = E(real_images).view(_batch_size,-1)\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),requires_grad=False).to(device=device)\n",
    "\n",
    "        x_hat = G(z_hat)\n",
    "        x_rand = G(z_rand)\n",
    "\n",
    "        l1_loss = 10 * criterion_l1(x_hat, real_images)\n",
    "        c_loss = criterion_bce(CD(z_hat), real_y[:_batch_size])\n",
    "        d_real_loss = criterion_bce(D(x_hat), real_y[:_batch_size]) \n",
    "        d_fake_loss = criterion_bce(D(x_rand), real_y[:_batch_size])\n",
    "\n",
    "        loss1 = l1_loss + c_loss + d_real_loss + d_fake_loss\n",
    "\n",
    "        loss1.backward(retain_graph=True)\n",
    "        e_optimizer.step()\n",
    "\n",
    "        g_optimizer.step()\n",
    "        g_optimizer.step()\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(1):\n",
    "        d_optimizer.zero_grad()\n",
    "\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),volatile=True).to(device=device)\n",
    "        z_hat = E(real_images).view(_batch_size,-1)\n",
    "        x_hat = G(z_hat)\n",
    "        x_rand = G(z_rand)\n",
    "\n",
    "        x_loss2 = 2.0 * criterion_bce(D(real_images), real_y[:_batch_size])+criterion_bce(D(x_hat), fake_y[:_batch_size])\n",
    "        z_loss2 = criterion_bce(D(x_rand), fake_y[:_batch_size])\n",
    "        loss2 = x_loss2 + z_loss2\n",
    "\n",
    "        if iters<4:\n",
    "            loss2.backward(retain_graph=True)\n",
    "        else:\n",
    "            loss2.backward(retain_graph=True)\n",
    "        d_optimizer.step()\n",
    "    ###############################################\n",
    "    # Train CD\n",
    "    ###############################################\n",
    "    for p in D.parameters():  # reset requires_grad\n",
    "        p.requires_grad = False\n",
    "    for p in CD.parameters():  # reset requires_grad\n",
    "        p.requires_grad = True\n",
    "    for p in E.parameters():  # reset requires_grad\n",
    "        p.requires_grad = False\n",
    "    for p in G.parameters():  # reset requires_grad\n",
    "        p.requires_grad = False\n",
    "\n",
    "    for iters in range(1):\n",
    "        cd_optimizer.zero_grad()\n",
    "        z_hat = E(real_images).view(_batch_size,-1)\n",
    "        x_loss3 = criterion_bce(CD(z_hat), fake_y[:_batch_size])\n",
    "        z_rand = Variable(torch.randn((_batch_size,latent_dim)),volatile=True).to(device=device)\n",
    "        z_loss3 = criterion_bce(CD(z_rand), real_y[:_batch_size])\n",
    "        loss3 = x_loss3 + z_loss3\n",
    "        loss3.backward(retain_graph=True)\n",
    "        cd_optimizer.step()\n",
    "        \n",
    "    ###############################################\n",
    "    # Visualization\n",
    "    ###############################################\n",
    "\n",
    "    if iteration % 100 == 0:\n",
    "        print('[{}/{}]'.format(iteration,MAX_ITER),\n",
    "              'D: {:<8.3}'.format(loss2.item()), \n",
    "              'En_Ge: {:<8.3}'.format(loss1.item()),\n",
    "              'Code: {:<8.3}'.format(loss3.item()))\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=\"Real\")\n",
    "        plotting.show()\n",
    "\n",
    "        featmask = np.squeeze((0.5*x_hat[0]+0.5).data.cpu().numpy())\n",
    "        featmask = nib.Nifti1Image(featmask,affine = np.eye(4))\n",
    "        plotting.plot_img(featmask,title=\"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=\"Rand\")\n",
    "        plotting.show()\n",
    "\n",
    "    if (iteration+1)%1000 == 0: \n",
    "        print(f'currrent iteration: {iteration}')\n",
    "        torch.save(G.state_dict(),'./checkpoint/G_noW_iter'+str(iteration+1)+'.pth')\n",
    "        torch.save(D.state_dict(),'./checkpoint/D_noW_iter'+str(iteration+1)+'.pth')\n",
    "        torch.save(E.state_dict(),'./checkpoint/E_noW_iter'+str(iteration+1)+'.pth')\n",
    "        torch.save(CD.state_dict(),'./checkpoint/CD_noW_iter'+str(iteration+1)+'.pth')\n",
    "        df = add_loss(df, iteration, loss1.item())\n",
    "        write_loss(df)\n",
    "        \n",
    "    iteration += 1\n"
   ]
  },
  {
   "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.6"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}