training_cv.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D U-Net on BraTS glioma dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext rpy2.ipython\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "import datetime\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import mxnet as mx\n",
    "from mxnet import gluon, autograd, ndarray as nd\n",
    "from mxnet.gluon.utils import split_and_load\n",
    "\n",
    "from unet import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## Setup hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = argparse.Namespace()\n",
    "\n",
    "args.data_dir = '../brats_2018_4D' # Should contain 'training' and 'validation' dirs. \n",
    "args.fold = 0 # Fold for 10-fold CV on training data. Comment out to include all data.\n",
    "\n",
    "# Training\n",
    "args.resume = ''\n",
    "args.start_epoch = 0\n",
    "args.epochs = 600\n",
    "args.batch_size = 6\n",
    "args.num_workers = 6\n",
    "# args.smooth = 1\n",
    "args.optimizer = 'adam'\n",
    "args.optimizer_params = {'learning_rate': 0.0001, 'lr_scheduler': mx.lr_scheduler.PolyScheduler(max_update=43*args.epochs, base_lr=0.0001, pwr=2)}\n",
    "# GPU_COUNT = 1\n",
    "# args.ctx = [mx.gpu(i) for i in range(GPU_COUNT)]\n",
    "args.ctx = [mx.gpu(0)]\n",
    "\n",
    "# Unet\n",
    "args.num_downs = 4 # Number of encoding/downsampling layers\n",
    "args.classes = 4 # Number of classes for segmentation, including background\n",
    "args.ngf = 32 # Number of channels in base/outermost layer\n",
    "args.use_bias = True # For conv blocks\n",
    "args.use_global_stats = True # For BN blocks\n",
    "\n",
    "# Pre/post-processing\n",
    "args.crop_size_train = [80, 80, 80] # Training patch size\n",
    "args.pad_size_val = [240, 240, 155] # Should be input vol dims unless 'crop_size_val' is larger\n",
    "args.crop_size_val = [192, 192, 144] # Should be divisible by 2^num_downs\n",
    "args.overlap = 0 # Fractional overlap for val patch prediction, combined with voting\n",
    "if hasattr(args, 'fold'): \n",
    "    args.fold_inds = get_k_folds(n=285, k=10, seed=1)[args.fold]\n",
    "args.lesion_frac = 0.9 # Fraction of patch centerpoints to be placed within lesion (vs randomly within brain)\n",
    "# args.warp_params = None\n",
    "args.warp_params = {'theta_max': 45,\n",
    "                    'offset_max': 0,\n",
    "                    'scale_max': 1.25,\n",
    "                    'shear_max': 0.1}\n",
    "\n",
    "# Checkpoint\n",
    "args.save_interval = args.epochs\n",
    "args.save_dir = '../params'\n",
    "args.val_interval = 10\n",
    "fold_str = 'fold' + str(args.fold) if hasattr(args, 'fold') else 'foldAll'\n",
    "time_str = str(datetime.datetime.now().strftime(\"%Y-%m-%dT%H%M%S\"))\n",
    "net_name = '_'.join(('unet', str(args.crop_size_train[0]), fold_str, time_str))\n",
    "args.save_prefix = os.path.join(args.save_dir, net_name, net_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## Setup data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.load('data/normalization_stats.npz')\n",
    "means_brain = nd.array(data['means_brain'])\n",
    "stds_brain  = nd.array(data['stds_brain'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = MRISegDataset4D(root=args.data_dir, split='train', mode='train', crop_size=args.crop_size_train, transform=brats_transform, means=means_brain, stds=stds_brain, fold_inds=args.fold_inds, lesion_frac=args.lesion_frac, warp_params=args.warp_params)\n",
    "valset   = MRISegDataset4D(root=args.data_dir, split='val',   mode='val',   crop_size=args.pad_size_val,    transform=brats_transform, means=means_brain, stds=stds_brain, fold_inds=args.fold_inds)\n",
    "\n",
    "train_data = gluon.data.DataLoader(trainset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True,  last_batch='rollover')\n",
    "val_data   = gluon.data.DataLoader(valset,   batch_size=1,               num_workers=args.num_workers, shuffle=False, last_batch='keep')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## Setup model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UnetGenerator(num_downs        = args.num_downs, \n",
    "                      classes          = args.classes, \n",
    "                      ngf              = args.ngf, \n",
    "                      use_bias         = args.use_bias, \n",
    "                      use_global_stats = args.use_global_stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.hybridize()\n",
    "model.collect_params().initialize(mx.init.Xavier(), force_reinit=True, ctx=args.ctx)\n",
    "if args.resume.strip():\n",
    "    model.load_parameters(os.path.join('../params', args.resume.strip()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = gluon.loss.SoftmaxCrossEntropyLoss(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = gluon.Trainer(model.collect_params(), args.optimizer, args.optimizer_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger, sw = start_logger(args)\n",
    "\n",
    "best_wt     = 0.\n",
    "best_avg    = 0.\n",
    "global_step = 0\n",
    "for epoch in range(args.start_epoch, args.epochs):\n",
    "    tbar = tqdm(train_data)\n",
    "    train_loss = 0.\n",
    "    for i, (data, label) in enumerate(tbar):\n",
    "        n_batch = data.shape[0]\n",
    "        label = label.squeeze(axis=1)\n",
    "        label = split_and_load(label, args.ctx)\n",
    "        data  = split_and_load(data,  args.ctx)\n",
    "        with autograd.record():\n",
    "            losses = [loss(model(X), Y) for X, Y in zip(data, label)]\n",
    "            for l in losses:\n",
    "                l.backward()\n",
    "                train_loss += l.sum().asnumpy() / n_batch\n",
    "        trainer.step(n_batch)\n",
    "\n",
    "        # Mini-batch logging\n",
    "        sw.add_scalar(tag='Cross_Entropy', value=('Train loss', l.mean().asscalar()), global_step=global_step)\n",
    "        global_step += 1\n",
    "        if i < (len(tbar) - 1):\n",
    "            tbar.set_description('E %d | loss %.4f'%(epoch, train_loss/(i+1)))\n",
    "    \n",
    "    # Epoch logging\n",
    "    if (epoch + 1) % args.val_interval == 0:\n",
    "        mx.nd.waitall()\n",
    "        metrics = brats_validate(model, val_data, crop_size=args.crop_size_val, overlap=args.overlap, ctx=args.ctx[0])\n",
    "        DSC_avg = log_epoch_hooks(epoch, train_loss/(i+1), metrics, logger, sw)\n",
    "        best_wt  = save_params(model, best_wt, metrics['WT']['DSC'], epoch, args.save_interval, args.save_prefix)\n",
    "        best_avg = save_params(model, best_avg, DSC_avg, epoch, None, args.save_prefix + '_avg')\n",
    "sw.export_scalars(args.save_prefix + '_scalar_dict.json')\n",
    "sw.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:unet]",
   "language": "python",
   "name": "conda-env-unet-py"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}