training_all_svid5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D U-Net on BraTS glioma dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "\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.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.lesion_frac = 0.9 # Fraction of patch centerpoints to be placed within lesion (vs randomly within brain)\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",
    "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",
    "## Overrides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args.classes = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MRISegDataset4DSVID(MRISegDataset):\n",
    "    \"\"\"Dataset class with all inputs and GT seg combined into a single 4D NifTI\"\"\"\n",
    "    def __getitem__(self, idx):\n",
    "        #import pdb; pdb.set_trace()\n",
    "        _sub_name = os.path.basename(os.path.dirname(self.sub_dirs[idx]))\n",
    "\n",
    "        # Load multichannel input data\n",
    "        img_path = os.path.join(self.sub_dirs[idx], _sub_name + '_' + '4D' + '.nii.gz')\n",
    "        img_raw = nib.load(img_path).get_fdata()\n",
    "        img_raw = img_raw.transpose((3,0,1,2))\n",
    "        img_raw = np.flip(img_raw, 2) # Correct AP orientation\n",
    "        img = img_raw[0:4]\n",
    "        \n",
    "        # Load segmentation label map\n",
    "        if self.split is not 'test':\n",
    "            target = img_raw[4]\n",
    "            #target[target==4] = 3 # Need to have consecutive integers [0, n_classes) for training\n",
    "            \n",
    "            # Include SVID=3 labels\n",
    "            seg_dir = 'data/BRATS_Training_WMH_Revised'\n",
    "            target_svid = nib.load(os.path.join(seg_dir, _sub_name + '_seg2.nii.gz')).get_fdata()\n",
    "            target_svid = np.flip(target_svid, 1) # Correct AP orientation\n",
    "            target[target_svid==2] = 3\n",
    "        else:\n",
    "            target = np.zeros_like(img[0,:]) # dummy segmentation\n",
    "        target = np.expand_dims(target, axis=0)\n",
    "\n",
    "        # Data augmentation\n",
    "        if self.mode == 'train':\n",
    "            img, target = self._sync_transform(img, target)\n",
    "        elif self.mode == 'val':\n",
    "            img, target = self._val_sync_transform(img, target)\n",
    "        else:\n",
    "            raise RuntimeError('unknown mode for dataloader: {}'.format(self.mode))\n",
    "\n",
    "        # Routine img specific processing (normalize, etc.)\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img, self.means[idx], self.stds[idx])\n",
    "\n",
    "        return img, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_brats_metrics():\n",
    "    \"\"\"Initialize dict for BraTS Dice metrics\"\"\"\n",
    "    metrics = {}\n",
    "    metrics['ET'] = {'labels': [4]}\n",
    "    metrics['TC'] = {'labels': [1, 4]}\n",
    "    metrics['WT'] = {'labels': [1, 2, 4]}\n",
    "    metrics['SVID'] = {'labels': [3]}\n",
    "    for _, value in metrics.items():\n",
    "        value.update({'tp':0, 'tot':0})\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def brats_validate(model, data_loader, crop_size, overlap, ctx):\n",
    "    \"\"\"Predict segs from val data, compare to ground truth val segs, and calculate val dice metrics\"\"\"\n",
    "    # Setup metric dictionary\n",
    "    metrics = init_brats_metrics()\n",
    "\n",
    "    # Get patch index iterator\n",
    "    dims = data_loader._dataset[0][1].shape[1:]\n",
    "    patch_iter = get_patch_iter(dims, crop_size, overlap)\n",
    "    \n",
    "    # Iterate over subjects\n",
    "    for i, (data, label) in enumerate(data_loader):  \n",
    "\n",
    "        # Iterate over patches\n",
    "        for inds in patch_iter:\n",
    "            data_patch  = get_patch(data, inds).as_in_context(ctx)\n",
    "            label_patch = get_patch(label, inds)\n",
    "            label_mask = label_patch.squeeze().asnumpy()\n",
    "            \n",
    "            output_mask = get_output_mask(model, data_patch).asnumpy()\n",
    "\n",
    "            # Update metrics\n",
    "            for _, metric in metrics.items():\n",
    "                label_mask_bin  = np.isin(label_mask, metric['labels'])\n",
    "                output_mask_bin = np.isin(output_mask, metric['labels'])\n",
    "                metric['tp']  += np.sum(label_mask_bin * output_mask_bin)\n",
    "                metric['tot'] += np.sum(label_mask_bin) + np.sum(output_mask_bin)\n",
    "\n",
    "    # Calculate overall metrics\n",
    "    for _, metric in metrics.items():\n",
    "            metric['DSC'] = 2 * metric['tp'] / metric['tot']\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_epoch_hooks(epoch, train_loss, metrics, logger, sw):\n",
    "    \"\"\"Epoch logging\"\"\"\n",
    "    DSCs = np.array([v['DSC'] for k,v in metrics.items()])\n",
    "    DSC_avg = DSCs.mean()\n",
    "    logger.info('E %d | loss %.4f | ET %.4f | TC %.4f | WT %.4f | SVID %.4f | Avg %.4f'%((epoch, train_loss) + tuple(DSCs) + (DSC_avg, )))\n",
    "    sw.add_scalar(tag='Dice', value=('Val ET', DSCs[0]), global_step=epoch)\n",
    "    sw.add_scalar(tag='Dice', value=('Val TC', DSCs[1]), global_step=epoch)\n",
    "    sw.add_scalar(tag='Dice', value=('Val WT', DSCs[2]), global_step=epoch)\n",
    "    sw.add_scalar(tag='Dice', value=('Val SVID', DSCs[3]), global_step=epoch)\n",
    "    sw.add_scalar(tag='Dice', value=('Val Avg', DSCs.mean()), global_step=epoch)\n",
    "    return DSC_avg"
   ]
  },
  {
   "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 = MRISegDataset4DSVID(root=args.data_dir, split='train', mode='train', crop_size=args.crop_size_train, transform=brats_transform, means=means_brain, stds=stds_brain, lesion_frac=args.lesion_frac, warp_params=args.warp_params)\n",
    "train_data = gluon.data.DataLoader(trainset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True,  last_batch='rollover')"
   ]
  },
  {
   "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",
    "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",
    "        tbar.set_description('E %d | loss %.4f'%(epoch, train_loss/(i+1)))\n",
    "    \n",
    "    # Epoch logging\n",
    "    best_wt  = save_params(model, 1, 0, epoch, args.save_interval, args.save_prefix)\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
}