MNIST Example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Learning Image Registration (DLIR) Framework\n",
    "## An example on MNIST\n",
    "\n",
    "TorchIR is a image registration library for **deep learning image registration (DLIR)**. I have \n",
    "integrated several ideas for image registration.\n",
    "\n",
    "I rely on PyTorch Lightning, which can be installed via:\n",
    "> pip install pytorch-lightning\n",
    "\n",
    "The pytorch-lightning trainer modules automatically create tensorboard log files. I store them in \n",
    "the `./output/lightning_logs` directory. Simply inspect them using:\n",
    "> tensorboard --logdir=./output/lightning_logs\n",
    "\n",
    "If you use this code for your publications, don't forget to cite my work ;)\n",
    "\n",
    "[1] Bob D. de Vos, Floris F. Berendsen, Max A. Viergever, Marius Staring and Ivana Išgum, \n",
    "\"End-to-end unsupervised deformable image registration with a convolutional neural network,\" \n",
    "in Deep learning in medical image analysis and multimodal learning for clinical decision support. \n",
    "Springer, Cham, 2017. p. 204-212, doi: 10.1007/978-3-319-67558-9_24\n",
    "https://link.springer.com/chapter/10.1007%2F978-3-319-67558-9_24\n",
    "\n",
    "[2] Bob D. de Vos, Floris F. Berendsen, Max A. Viergever, Hessam Sokooti, Marius Staring and Ivana Išgum\n",
    "\"A deep learning framework for unsupervised affine and deformable image registration,\" Medical image analysis, vol. 52, pp. 128-143, Feb. 2019, doi: 10.1016/j.media.2018.11.010\n",
    "https://www.sciencedirect.com/science/article/pii/S1361841518300495\n",
    "\n",
    "Please note that the code is still under heavy development and I'd really love your input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "\n",
    "%autoreload 2\n",
    "    \n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from pathlib import Path\n",
    "DEST_DIR = Path('./output')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The data\n",
    "To illustrate the workings of the code I'll set up a registration experiment using MNIST data. I also used these in my first publication on deep learning image registration [1]. We will setup the experiments such that only one number will be loaded and we will register those to eachother.\n",
    "\n",
    "We'll start with loading some modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from torchir.utils import IRDataSet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we create a specific dataset class that selects MNIST instances with a specific class label. Here we take all number 9s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTSubSet(MNIST):\n",
    "    '''\n",
    "    A Dataset class that selects a single type of MNIST digit.\n",
    "    '''\n",
    "    def __init__(self, label, rng=np.random.default_rng(), *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        assert(label >= 0 and label <= 9)\n",
    "        idcs = torch.where(self.targets == label)\n",
    "        \n",
    "        self.data = self.data[idcs]\n",
    "        self.targets = self.targets[idcs]\n",
    "\n",
    "        self.transform = transform\n",
    "        self.rng = rng\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        return super().__getitem__(idx)[0] # only return image\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we set up the dataloaders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(808)\n",
    "transform=transforms.Compose([transforms.ToTensor(),\n",
    "                              transforms.Normalize((0.5), (0.5)),\n",
    "                             ])\n",
    "ds_train_subset = MNISTSubSet(label=9, rng=rng, root='../datasets/',  transform=transform, download=True, train=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We reserve part of the training set for validation. Note that we only select a small number for validation. This will become clear immediately after this block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_set_size = 20\n",
    "train_set_size = len(ds_train_subset) - val_set_size\n",
    "ds_train_subset, ds_validation_subset = torch.utils.data.random_split(ds_train_subset, [train_set_size, val_set_size], \n",
    "                                                        generator=torch.Generator().manual_seed(808))\n",
    "ds_test_subset = MNISTSubSet(label=9, rng=rng, root='../datasets/',  transform=transform, download=True, train=False)\n",
    "print(f'Training subset size: {len(ds_train_subset)}')\n",
    "print(f'Validation subset size: {len(ds_validation_subset)}')\n",
    "print(f'Test subset size: {len(ds_test_subset)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will employ a convenience class to convert the training and validation data sets into image registration sets. The new class will provide all possible permutations of the input dataset. Note that this will heavily increase the number of training instances (i.e. there are many fixed and moving image permutations). Note that we will use the test set later in a different manner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train = IRDataSet(ds_train_subset)\n",
    "ds_validation = IRDataSet(ds_validation_subset)\n",
    "print(f'Training IR set size: {len(ds_train)}')\n",
    "print(f'Validation IR set size: {len(ds_validation)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each instance of an IR dataset is a permutation of a fixed and moving image. This results in a very high number of possible permutations. Since Pytorch Lightning does not play well with iterations and rather likes epochs, we limit the number of permutations per epoch by setting up our own data samplers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "training_batches = 100\n",
    "validation_batches = 10\n",
    "\n",
    "train_sampler = torch.utils.data.RandomSampler(ds_train, replacement=True, \n",
    "                                               num_samples=training_batches*batch_size, \n",
    "                                               generator=torch.Generator().manual_seed(808))\n",
    "train_loader = torch.utils.data.DataLoader(ds_train, batch_size, sampler=train_sampler)\n",
    "val_loader = torch.utils.data.DataLoader(ds_validation, batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have setup all our data classes and now we can start the image registration experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single DIRNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will first start with deformable image registration using a single layer of b-spline registration. For this, we will initialize a DIRNet model, which is a straight-forward network [1, 2]. The DIRNet takes two images, a fixed and moving image, as its input and it outputs a number of control points: output shape is: num_samples, ndim=2, image_height//grid_spacing, image_width//grid_spacing. \n",
    "\n",
    "The DIRNet is initialized with a grid spacing of (8, 8) voxels. The network can handle a varying input of image sizes, because the network is implemented in a fully convolutional way. The order of the b-spline is set at 3 (i.e. cubic), which is common for Bspline registration.\n",
    "\n",
    "We use the model to initialize a BsplineTransformer layer. The BsplineTransformer is tasked with interpolating the b-spline control points to a full displacement vector field (DVF). For convenience, the BsplineTransformer will automatically resample the moving image, using the generated DVF.\n",
    "\n",
    "To simplify the code, I have used pytorch-lighting. Pytorch lighthning can easily be installed via `pip install pytorch-lighting`\n",
    "\n",
    "Pytorch lightning may look a bit daunting if you're not used to it, but I am sure you'll learn to appreciate it when you notice how it nicely procedures things such as logging and pushing data to the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl\n",
    "from torchir.networks import DIRNet\n",
    "from torchir.metrics import NCC\n",
    "from torchir.transformers import BsplineTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LitDIRNet(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        grid_spacing = (8, 8)\n",
    "        self.dirnet = DIRNet(kernels=16, grid_spacing=grid_spacing)\n",
    "        self.bspline_transformer = BsplineTransformer(ndim=2, upsampling_factors=grid_spacing)\n",
    "        self.metric = NCC()\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        lr = 0.001\n",
    "        optimizer = torch.optim.Adam(self.dirnet.parameters(), lr=lr, amsgrad=True)\n",
    "        return optimizer\n",
    "\n",
    "    def forward(self, fixed, moving):\n",
    "        params = self.dirnet(fixed, moving)\n",
    "        warped = self.bspline_transformer(params, fixed, moving)\n",
    "        return warped\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/training', loss)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/validation', loss)\n",
    "        return loss  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we set up the trainer and start training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = LitDIRNet()\n",
    "trainer = pl.Trainer(default_root_dir=DEST_DIR, \n",
    "                     log_every_n_steps=50,\n",
    "                     val_check_interval=50, \n",
    "                     max_epochs=100, \n",
    "                     gpus=1)\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We save the network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_checkpoint(DEST_DIR / 'mnist_ir_8.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To illustrate its performance, we recreate an experiment from my paper [1] where we take one fixed image from the test set and register all other images to that target image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_moving = np.zeros((28, 28), dtype=float)\n",
    "avg_warped = np.zeros((28, 28), dtype=float)\n",
    "model = model.cuda()\n",
    "model.eval()\n",
    "fixed = ds_test_subset[0]\n",
    "with torch.no_grad():\n",
    "    for moving in tqdm(ds_test_subset):\n",
    "        warped = model(fixed[None].cuda(), moving[None].cuda()).detach().squeeze().cpu().numpy()\n",
    "        avg_moving += moving.squeeze().cpu().numpy() / len(ds_test_subset)\n",
    "        avg_warped += warped / len(ds_test_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(1, 3, figsize=(10, 3.5))\n",
    "axarr[0].imshow(-fixed.squeeze().cpu().numpy(), cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[1].imshow(-avg_moving, cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[2].imshow(-avg_warped, cmap='gray', vmin=-1, vmax=1)\n",
    "for ax in axarr.ravel():\n",
    "    ax.axis('off')\n",
    "\n",
    "axarr[0].set_title('fixed image')\n",
    "axarr[1].set_title('images before registration')\n",
    "axarr[2].set_title('images after registration')\n",
    "\n",
    "fig.suptitle('Deformable image registration network (DIRNet)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to inspect the data, use tensorboard. The logs are stored in the specified `DEST_DIR`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Affine registration\n",
    "\n",
    "We can also do affine registration, which I showed in [2].\n",
    "Note that with the data we use here, the fixed and moving batches are of similar size, but, as demonstrated in the paper, they can be of different size. The AIRNet allows this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchir.networks import AIRNet\n",
    "from torchir.transformers import AffineTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LitAIRNet(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.airnet = AIRNet(kernels=16)\n",
    "        self.global_transformer = AffineTransformer(ndim=2)\n",
    "        self.metric = NCC()\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        lr = 0.001\n",
    "        optimizer = torch.optim.Adam(self.airnet.parameters(), lr=lr, amsgrad=True)\n",
    "        return optimizer\n",
    "\n",
    "    def forward(self, fixed, moving):\n",
    "        parameters = self.airnet(fixed, moving)\n",
    "        warped  = self.global_transformer(parameters, fixed, moving)\n",
    "        return warped\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/training', loss)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/validation', loss)\n",
    "        return loss  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = LitAIRNet()\n",
    "trainer = pl.Trainer(default_root_dir=DEST_DIR,\n",
    "                     log_every_n_steps=50,\n",
    "                     val_check_interval=50,\n",
    "                     max_epochs=100,\n",
    "                     gpus=1)\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We save the network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_checkpoint(DEST_DIR / 'mnist_ir_affine.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we check if it can register images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_moving = np.zeros((28, 28), dtype=float)\n",
    "avg_warped = np.zeros((28, 28), dtype=float)\n",
    "model = model.cuda()\n",
    "fixed = ds_test_subset[0]\n",
    "for moving in tqdm(ds_test_subset):\n",
    "    warped = model(fixed[None].cuda(), moving[None].cuda()).detach().squeeze().cpu().numpy()\n",
    "    avg_moving += moving.squeeze().cpu().numpy() / len(ds_test_subset)\n",
    "    avg_warped += warped / len(ds_test_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(1, 3, figsize=(10, 3.5))\n",
    "axarr[0].imshow(-fixed.squeeze().cpu().numpy(), cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[1].imshow(-avg_moving, cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[2].imshow(-avg_warped, cmap='gray', vmin=-1, vmax=1)\n",
    "for ax in axarr.ravel():\n",
    "    ax.axis('off')\n",
    "\n",
    "axarr[0].set_title('fixed image')\n",
    "axarr[1].set_title('images before registration')\n",
    "axarr[2].set_title('images after registration')\n",
    "\n",
    "fig.suptitle('Affine image registration network (AIRNet)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DLIR Framework"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In [2] I also demonstrated that multiple coarse-to-fine registration layers improve image registration. Now let's implement this using the DLIRFramework module. I chose a dynamic implementation where we add a layer. We train the layer. We add another layer and fix the weight of the previous layer. We train the new layer. Etc..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchir.dlir_framework import DLIRFramework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LitDLIRFramework(pl.LightningModule):\n",
    "    def __init__(self, only_last_trainable=True):\n",
    "        super().__init__()\n",
    "        self.dlir_framework = DLIRFramework(only_last_trainable=only_last_trainable)\n",
    "        self.add_stage = self.dlir_framework.add_stage\n",
    "        self.metric = NCC()\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        lr = 0.001\n",
    "        weight_decay = 0\n",
    "        optimizer = torch.optim.Adam(self.dlir_framework.parameters(), lr=lr, weight_decay=weight_decay, amsgrad=True)\n",
    "        return {'optimizer': optimizer}\n",
    "\n",
    "    def forward(self, fixed, moving):\n",
    "        warped = self.dlir_framework(fixed, moving)\n",
    "        return warped\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/training', loss)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        warped = self(batch['fixed'], batch['moving'])\n",
    "        loss = self.metric(batch['fixed'], warped)\n",
    "        self.log('NCC/validation', loss)\n",
    "        return loss  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the model, and add an affine registration layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = LitDLIRFramework()\n",
    "model.add_stage(network=AIRNet(kernels=16), transformer=AffineTransformer(ndim=2))\n",
    "trainer = pl.Trainer(default_root_dir=DEST_DIR,\n",
    "                     log_every_n_steps=50,\n",
    "                     val_check_interval=50,\n",
    "                     max_epochs=100,\n",
    "                     gpus=1)\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add a BSpline layer with an 8x8 grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.add_stage(network=DIRNet(grid_spacing=(8, 8), kernels=16, num_conv_layers=5, num_dense_layers=2),\n",
    "                transformer=BsplineTransformer(ndim=2, upsampling_factors=(8, 8)))\n",
    "trainer = pl.Trainer(default_root_dir=DEST_DIR,\n",
    "                     log_every_n_steps=50,\n",
    "                     val_check_interval=50,\n",
    "                     max_epochs=100,\n",
    "                     gpus=1)\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add a finer 4x4 grid and train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.add_stage(network=DIRNet(grid_spacing=(4, 4), kernels=16, num_conv_layers=5, num_dense_layers=2),\n",
    "                transformer=BsplineTransformer(ndim=2, upsampling_factors=(4, 4)))\n",
    "trainer = pl.Trainer(default_root_dir=DEST_DIR,\n",
    "                     log_every_n_steps=50,\n",
    "                     val_check_interval=50,\n",
    "                     max_epochs=100,\n",
    "                     gpus=1)\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's store the checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_checkpoint(DEST_DIR / 'mnist_dlir_affine_8_4.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now do ultra-fast coarse-to-fine image registration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_moving = np.zeros((28, 28), dtype=float)\n",
    "avg_warped = np.zeros((28, 28), dtype=float)\n",
    "model = model.cuda()\n",
    "fixed = ds_test_subset[0]\n",
    "for moving in tqdm(ds_test_subset):\n",
    "    warped = model(fixed[None].cuda(), moving[None].cuda()).detach().squeeze().cpu().numpy()\n",
    "    avg_moving += moving.squeeze().cpu().numpy() / len(ds_test_subset)\n",
    "    avg_warped += warped / len(ds_test_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axarr = plt.subplots(1, 3, figsize=(10, 3.5))\n",
    "axarr[0].imshow(-fixed.squeeze().cpu().numpy(), cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[1].imshow(-avg_moving, cmap='gray', vmin=-1, vmax=1)\n",
    "axarr[2].imshow(-avg_warped, cmap='gray', vmin=-1, vmax=1)\n",
    "for ax in axarr.ravel():\n",
    "    ax.axis('off')\n",
    "\n",
    "axarr[0].set_title('fixed image')\n",
    "axarr[1].set_title('images before registration')\n",
    "axarr[2].set_title('images after registration')\n",
    "\n",
    "fig.suptitle('DLIR Framework: coarse-to-fine b-splines');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading a model for later use can be done as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LitDLIRFramework()\n",
    "model.add_stage(GlobalTransformer(AIRNet(kernels=16)))\n",
    "model.add_stage(BsplineTransformer(DIRNet(grid_spacing=(8, 8), kernels=16, num_conv_layers=5, num_dense_layers=2)))\n",
    "model.add_stage(BsplineTransformer(DIRNet(grid_spacing=(4, 4), kernels=16, num_conv_layers=5, num_dense_layers=2)))\n",
    "model.load_state_dict(torch.load(DEST_DIR / 'mnist_dlir_affine_8_4.ckpt')['state_dict'])"
   ]
  }
 ],
 "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"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}