VAEGAN_ADNI_GP_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",
    "\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 *\n",
    "from Model_alphaWGAN import Discriminator as alpha_D\n",
    "from utils import *\n",
    "from utils import sinkhorn_pointcloud as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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\n",
    "\n",
    "gpu_0 = 0\n",
    "gpu_1 = 0\n",
    "torch_seed = 4\n",
    "r_g = torch.manual_seed(torch_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Encoder(\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",
       "  (mean): Sequential(\n",
       "    (0): Linear(in_features=32768, out_features=2048, bias=True)\n",
       "    (1): BatchNorm1d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): Linear(in_features=2048, out_features=1000, bias=True)\n",
       "  )\n",
       "  (logvar): Sequential(\n",
       "    (0): Linear(in_features=32768, out_features=2048, bias=True)\n",
       "    (1): BatchNorm1d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): Linear(in_features=2048, out_features=1000, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "G = Generator(noise = latent_dim)\n",
    "D = Discriminator()\n",
    "E = Encoder(gpu_ind=gpu_0)\n",
    "# E = alpha_D(out_class = latent_dim, is_dis=False)\n",
    "\n",
    "G.cuda(gpu_0)\n",
    "D.cuda(gpu_0)\n",
    "E.cuda(gpu_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCH = 300\n",
    "N_ITER = 100000\n",
    "\n",
    "# criterion_bce = nn.BCELoss()\n",
    "# criterion_l1 = nn.L1Loss()\n",
    "\n",
    "d_iter = 6\n",
    "\n",
    "pth = 'vae_checkpoint_gp'\n",
    "# load the highest savepoints of all models\n",
    "df = load_loss(f'./{pth}/loss.csv')\n",
    "# init_epoch = vae_load_checkpoint(G, D, E)\n",
    "init_iter = load_checkpoint(G, D, E, None, '_iter', path=pth) + 2\n",
    "gen_load = inf_train_gen(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/projects/shibo/3dbraingen/utils/__init__.py:27: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  yield torch.tensor(images, requires_grad=False)\n",
      "/usr/local/miniconda3/lib/python3.7/site-packages/torch/nn/functional.py:2351: UserWarning: nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.\n",
      "  warnings.warn(\"nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.\")\n",
      "/usr/local/miniconda3/lib/python3.7/site-packages/torch/nn/functional.py:1320: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.\n",
      "  warnings.warn(\"nn.functional.tanh is deprecated. Use torch.tanh instead.\")\n",
      "/usr/local/miniconda3/lib/python3.7/site-packages/torch/nn/functional.py:1332: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n"
     ]
    }
   ],
   "source": [
    "# for epoch in range(init_epoch, N_EPOCH):\n",
    "for iteration in range(init_iter, N_ITER):\n",
    "#         real_y = torch.ones((BATCH_SIZE, 1), requires_grad=False)\n",
    "#         fake_y = torch.zeros((BATCH_SIZE, 1), requires_grad=False)\n",
    "\n",
    "    real_images = gen_load.__next__().cuda(gpu_0)\n",
    "    _batch_size = real_images.size(0)\n",
    "    z_rand = torch.randn((_batch_size, latent_dim), requires_grad=False).cuda(gpu_0)\n",
    "\n",
    "    mean,logvar,code = E(real_images)\n",
    "    x_hat = G(code)\n",
    "    x_rand = G(z_rand)\n",
    "    d_real_loss = -D(real_images).mean()\n",
    "    d_recon_loss = D(x_hat).mean()\n",
    "    d_fake_loss = D(x_rand).mean()\n",
    "    gen_img_loss = -(d_real_loss + d_recon_loss+ d_fake_loss)\n",
    "    ###############################################\n",
    "    # Train D \n",
    "    ###############################################\n",
    "    for i in range(d_iter):\n",
    "        d_optimizer.zero_grad()\n",
    "\n",
    "        gradient_penalty_h = calc_gradient_penalty(D,real_images, x_hat, cuda_ind=gpu_0)\n",
    "        gradient_penalty_r = calc_gradient_penalty(D,real_images, x_rand, cuda_ind=gpu_0)\n",
    "\n",
    "        dis_loss = d_recon_loss + d_real_loss + d_fake_loss + gradient_penalty_r + gradient_penalty_h\n",
    "        dis_loss.backward(retain_graph=True)\n",
    "        d_optimizer.step()\n",
    "\n",
    "    # save memory\n",
    "    gradient_penalty_r = None\n",
    "    gradient_penalty_h = None \n",
    "\n",
    "    ###############################################\n",
    "    # Train G\n",
    "    ###############################################\n",
    "    g_optimizer.zero_grad()\n",
    "\n",
    "    rec_loss = ((x_hat - real_images)**2).mean()\n",
    "\n",
    "    err_dec = gamma * rec_loss + gen_img_loss\n",
    "    err_dec.backward(retain_graph=True)\n",
    "    g_optimizer.step()\n",
    "    ###############################################\n",
    "    # Train E KL\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",
    "    ## Train E WL ##############################\n",
    "#         z_hat = E(real_images).view(_batch_size,-1).cuda(gpu_0)\n",
    "#         code = z_hat\n",
    "#         w_dist = sp.sinkhorn_loss(torch.transpose(z_rand, 0, 1), torch.transpose(z_hat, 0, 1), 0.1, 1000, 100, gpu=gpu_0)\n",
    "#         err_enc = w_dist + beta*rec_loss\n",
    "    ###############################################\n",
    "\n",
    "    e_optimizer.zero_grad()\n",
    "    err_enc.backward()\n",
    "    e_optimizer.step()\n",
    "    ###############################################\n",
    "    # Visualization\n",
    "    ###############################################\n",
    "    if iteration % 500 == 0:\n",
    "        print(f'G loss: {err_dec.item()}')\n",
    "        print('[{}/{}]'.format(iteration, N_ITER),\n",
    "              'D: {:<8.3}'.format(dis_loss.item()), \n",
    "              'En: {:<8.3}'.format(err_enc.item()),\n",
    "              'De: {:<8.3}'.format(err_dec.item()) \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*G(code)[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*G(z_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",
    "    loss_dict = {\n",
    "        'd_real_loss': [d_real_loss.item()],\n",
    "        'd_recon': [d_recon_loss.item()],\n",
    "        'd_fake': [d_fake_loss.item()],\n",
    "        'err_enc': [err_enc.item()],\n",
    "    }\n",
    "\n",
    "\n",
    "    if iteration and iteration % 1000 == 0:\n",
    "        viz_pca(E, trainset, latent_size=latent_dim, index=iteration, is_cd=True, gpu_ind=gpu_0, batch_size=4)\n",
    "        viz_pca(G, trainset, latent_size=latent_dim, index=iteration, gpu_ind=gpu_0, batch_size=4)\n",
    "        torch.save(G.state_dict(),f'./{pth}/G_iter'+str(iteration)+'.pth')\n",
    "        torch.save(D.state_dict(),f'./{pth}/D_iter'+str(iteration)+'.pth')\n",
    "        torch.save(E.state_dict(),f'./{pth}/E_iter'+str(iteration)+'.pth')\n",
    "        df = add_loss(df, loss_dict)\n",
    "        write_loss(df, path=pth)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.unsqueeze(torch.randn((64, 64, 64)), 1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}