AirLab: Autograd Image Registration Laboratory
Medical image registration is an active research topic and forms a basis for many medical image analysis tasks. Although image registration is a rather general concept specialized methods are usually required to target a specific registration problem. The development and implementation of such methods has been tough so far as the gradient of the objective has to be computed. Also, its evaluation has to be performed preferably on a GPU for larger images and for more complex transformation models and regularization terms. This hinders researchers from rapid prototyping and poses hurdles to reproduce research results. There is a clear need for an environment which hides this complexity to put the modeling and the experimental exploration of registration methods into the foreground. With the “Autograd Image Registration Laboratory” (AirLab), we introduce an open laboratory for image registration tasks, where the analytic gradients of the objective function are computed automatically and the device where the computations are performed, on a CPU or a GPU, is transparent. It is meant as a laboratory for researchers and developers enabling them to rapidly try out new ideas for registering images and to reproduce registration results which have already been published. AirLab is implemented in Python using PyTorch as tensor and optimization library and SimpleITK for basic image IO. Therefore, it profits from recent advances made by the machine learning community concerning optimization and deep neural network models.
The present draft of this paper roughly outlines AirLab with first code snippets and performance analyses. A more exhaustive introduction will follow as a final version soon.
The registration of images is a growing research topic and forms an integral part in many medical image analysis tasks . It is referred to as the process of finding corresponding structures within different images. There is a large number of applications where image registration is inevitable such as e.g. the fusion of different modalities, monitoring anatomical changes, population modelling or motion extraction.
Image registration is a nonlinear, ill-posed problem which is approached by optimizing a regularized objective. What is defined as quite general requires usually specialized objective functionals and implementations for applying it to specific registration tasks. The development of such specific registration methods has been tough so far and their implementation tedious. This is because gradients have to be computed within the optimization whose implementations are error-prone, especially for 3D objectives. Furthermore, for large 3D images, the computational demand is usually high and a parallel execution on a GPU unavoidable. These are problems which hinder researchers from playing around with different combinations of objectives and regularizers and rapidly trying out new ideas. Similarly, the effort to reproduce registration results is often out of proportion. There is a clear need for an environment which hides this complexity, enables rapid prototyping and simplifys reproduction.
In this paper, we introduce “Autograd Image Registration Laboratory” (AirLab), an image registration environment - or a laboratory - for rapid prototyping and reproduction of image registration methods. Thus, it addresses researchers and developers and simplifies their work in the exploration of different registration methods, in particular also with upcoming complex deep learning approaches. It is written in Python and based on the tensor library PyTorch . It heavily uses features from PyTorch such as autograd, the rich family of optimizers and the transparent utilization of GPUs. In addition, SimpleITK  is included for data input/output to support all standard image file formats. AirLab comes along with state-of-the-art registration components including various image similarity measures, regularization terms and optimizers. Experimenting with such building blocks or trying out new ideas, for say a regularizer, becomes almost effortless as gradients are computed automatically. Finally, example implementations of standard image registration methods are provided such as Optical Flow , Demons  and Free Form Deformations . AirLab is licensed under the Apache License 2.0 and available on GitHub111https://github.com/airlab-unibas/airlab.
In the following, we first provide a brief background about medical image registration followed by the description of AirLab, its building blocks and its features. Finally, we provide registration experiments with standard registration methods which are implemented in AirLab including performance analyses and code snippets.
The present draft of this paper roughly introduces AirLab and is intended for the presentation at the 8th International Workshop on Biomedical Image Registration in Leiden. A more detailed final version will follow soon.
Ii-a Image Registration
Let be a set of points arranged in a regular grid which covers the joint image domain of a moving and fixed image . The images map the -dimensional spatial domain to intensity values. Furthermore, let spatially transform the coordinate system of the moving image. Image registration can be formulated as a regularized minimization problem
where is a similarity measure between the transformed moving image and the fixed image and is a regularization term which operates on on the domain . The two terms are balanced by and is the function composition. An example for a similarity measure is the mean squared error measure for monomodal image registration
where is the cardinality of a set. An exemplary regularization term is the diffusion regularization which favours smooth transformations
Transformation models can be divided in basically four types: linear, non-linear/parametric, non-parametric and hybrid. Linear transformations, available in AirLab, transform each point with a linear map
where is an rotation/translation matrix up to in 2D and degrees of freedom in 3D and stands in homogeneous coordinates .
Non-linear/parametric transformations are defined in an interpolating basis on a coarse grid of control points
where and is the basis function. Common basis functions are the B-spline or Wendland kernel which both are implemented as example basis in AirLab (see Section III-B). The advantage of non-linear/parametric transformation models are, that they are computationally efficient. Furthermore, if is smooth they inherently yield smooth transformations.
In non-parametric methods, each point in the image can be transformed individually in dimensions giving a maximum flexibility. To still be able to reach reasonable registration results the regularization term is inevitable. Hierarchical models can be seen as hybrid parametric and non-parametric models. Their hierarchical structure enables them to capture large deformations . Non-parametric models are supported as well by AirLab while hybrid are planned.
The similarity measure depends nonlinearly on which makes an analytical solution to Equation 1 intractable. Because in non-linear registration the number of parameters of is in the millions, gradient based optimization is usually the only choice to reach a locally optimal transformation . Having PyTorch at hand, state-of-the-art optimizers are available in AirLab such as LBFGS  and Adam .
Ii-B Image Registration Frameworks
There are already a considerable amount of medical image registration frameworks available which are valuable enrichments to the community. Their focus and intensions are diverse, ranging from rich frameworks to specific implementations. For an exhaustive list and comparison of such image registration software we refer to . Gradient free approaches as e.g. the MRF-based method of  are out of scope of the current implementation of AirLab.
The Insight Segmentation and Registration Toolkit (ITK)  is a comprehensive framework for diverse image processing and analysis task written in C++. It is mainly intended for the use as a library for developers who want to implement ready-to-use software. The registration tool Elastix  is based on ITK and provides a collection of algorithms commonly used in image registration. It can also be used out-of-the-box with the possibility of a detailed configuration script. Furthermore, its plug-in architecture allows to integrate custom parts into the software. The extension SimpleElastix  offers bindings to other languages such as Python, Java, Ruby and more. Elastix and SimpleElastix are strong if one needs some flexibility in choosing and combining different registration components for a specific registration task. Scalable Image Analysis and Shape Modelling (Scalismo) [3, 19] is a library mainly for statistical shape modeling written in scala. It provides also image registration functionality and can be interactively executed similar to SimpleElastix. Advanced Normalization Tools (ANTs)  is based on ITK as well. It provides a command line tool including large deformation registration algorithms with standard similarity measures. The Automated Image Registration software AIR  is written in C and provides basic registration functionality for linear and polynomial non-linear image alignment up to the twelfth order. The Medical Image Registration ToolKit (MIRTK) [25, 27] is a collection of libraries and command-line tools for image and point-set registration. Various registration methods based on free form deformations are provided. Flexible Algorithms for Image Registration (FAIR)  is a software package written in MATLAB comprising various similarity measures and regularizers.
None of the mentioned software packages are suited for rapid prototyping in the development of image registration algorithms. This is mainly because: (I) For the optimization, gradients have to be provided. For complex transformation models, regularization terms and similarity measures, their implementation is highly error-prone. (II) For medical images, the computational demand is usually high and therefore the execution has to be performed on a GPU. The development for GPUs without an appropriate framework is not trivial. (III) The majority of the frameworks are written in C++. Thus, the development within those frameworks needs good expertise in this language. Furthermore, the number of code lines required for C++ implementations in these frameworks do not agree with the concept of rapid prototyping.
Iii Autograd Image Registration Laboratory
AirLab is a rapid prototyping environment for medical image registration. Its unique characteristics are the automatic differentiation and the transparent usage of GPUs. It is written in the scripting language Python and heavily uses key functionality of PyTorch .
The main building blocks constitute:
Iii-a Automatic Symbolic Differentiation
A key feature of AirLab is its automatic symbolic differentiation of the objective function. This means, that only the forward function has to be provided by the developer and the gradient which is required for the optimization is derived through automatic differentiation (AD). AirLab borrows the AD functionality of PyTorch. It is one of the fastest dynamic AD frameworks currently available. Its strong points are:
Dynamic: the function which is symbolically differentiated is defined by the computations which are run on the variables. Hence, no static graph structure has to be built which fosters rapid prototyping.
Immediate: only tensor computations which are necessary for differentiation are recorded
Core logic: a low overhead is needed as the AD logic is written in C++ and was carefully tuned
Please cf.  for more details.
Iii-B Image Registration Features
We list here the main building blocks required for medical image registration which are provided by AirLab. Upcoming features which are planed for implementation can be found in Section III-C.
Iii-B1 Similarly Measures
Mean Squared Errors: a simple and fast to compute point-wise measure which is well suited for monomodal image registration
Class name: MSELoss
Normalized Correlation Coefficient: a point-wise measure as . It is targeted to image registration tasks, where the intensity relation between the moving and the fixed images is linear
where the sums go over the image domain , E is the expectation value (or mean) and Var is the variance of the respective image.
Class name: NCCLoss
Local Cross Correlation: is the localized version of where the expectation value and the variance for a given are computed in a local neighborhood of . In AirLab is implemented with efficient convolution operations. Notice that the exact gradient is computed using autograd and no gradient approximation is performed in contrast to .
Class name: LCCLoss
Iii-B2 Transformation Models
AirLab supports three major types of transformation models: linear/parametric, non-linear/parametric and non-parametric models (hybrid models are planned).
Currently, AirLab supports rigid transformations including rotation and translation. Similarity and affine transformations are planned.
Class name: RigidTransformation
as mentioned with Equation 5, non-linear/parametric models have fewer control points as image points are available. The displacement for a given point in the image is interpolated from neighboring control points by the respective basis function. In AirLab, two exemplary basis functions are implemented:
B-spline: the standard B-spline kernel, which is used in the Free Form Deformation (FFD) algorithm of 
In addition, AirLab supports B-spline kernels of arbitrary order (first order are used in  and third order in the FFD ). An order is derived by convolving the zeroth order B-spline times with it self:
where corresponds to and is the convolution. The control points have a spacing of which implicitly defines the extent of the kernel. With increasing order, the control point support of the kernel is increased by one for each additional order.
Class name: BsplineTransformation
where and is the Wendland function of the second kind and positive definite in dimensions. The scaling can also be provided for each space dimension separately to achieve an anisotropic support.
Class name: WendlandKernelTransformation
For the non-linear/parametric transformation models, the transposed convolution is applied (cf. ) which is available in PyTorch. It is an up-sampling operation where the interpolation kernel can be provided. That means in our case, the control points are “up-sampled” and interpolated using the basis function of choice. Thus, their evaluation is highly optimized for CPUs and GPUs and the gradient computation, automatically performed with autograd.
the simpler model is the non-parametric model, where each point in the image can be independently transformed. That means, there are parameters (number of image points times number of space dimensions). To achieve a meaningful transformation, strong regularization is required.
Iii-B3 Image Warping
To compare the transformed moving image with the fixed image within the similarity measures the coordinate system of the moving image has to be warped. As it is mostly done in image registration, AirLab performs backward warping. That means, the transformation is defined on the fixed image domain where the displacement vectors point to the corresponding points in the moving image. To transform the moving image, it is backward warped into the coordinate system of the fixed image. This prevents holes occuring in the warped image.
The warping is performed in normalized coordiantes in the interval . The points which are transformed out of the fixed image region are identified by checking if falls outside the normalized interval. For illustration please see following snippet:
Because displaced points not necessarily fall onto the pixel-grid, interpolation is required. Currently, AirLab supports linear interpolation while B-spline interpolation is planned as up-coming feature. The warping is performed by the grid sampler of PyTorch which utilizes the GPU.
There are three different types of regularization terms in AirLab. (I) Regularizers on the displacement field , (II) regularizers on the parameters of and (III) the Demons regularizers which regularize the displacement field by filtering it in each iteration. Note that Demons regularizers are not differentiated, because in Demons approaches the optimization is an iteration scheme where the image forces (gradient of similarity measure) are evaluated to update the current displacement field and alternatingly the displacement field is regularized using filtering.
we first list the regularization terms which operate on the displacement field .
Diffusion: a regularizer which penalizes changes in the transformation
Class name: DiffusionRegulariser
Anisotropic Total Variation: a regularizer which favours piece-wise smooth transformations
It is anisotropic which means its influence is alinged to the coordinate axes.
Class name: TVRegulariser
Isotropic Total Variation: the isotropic version of the anisotropic regularizer
Both TV regularizers are not differentiable, therefore, the subgradient of zero is taken at zero.
Class name: IsotropicTVRegulariser
Sparsity: a regularizer which penalizes non-zero parameters
Class name: SparsityRegulariser
Regularizers on Parameters
The listed regularization terms are also available for regularizing the parameters of . The parameters which should be regularized are passed to the regularizer as an array, a name and a weighting. In this way, one can individually weight subsets of parameters, belonging for example to different hierarchical levels, cf. the following example:
Currently, there are two Demons regularizers available in AirLab:
Kernel: an arbitrary convolution kernel for filtering the displacement field. An example is the Gaussian kernel which is used originally in the Demons algorithm .
Class name: GaussianRegulariser
Graph Diffusion: the diffusion is performed by spectral graph diffusion. The graph can be utilized in order to handle the sliding organ problem. In this case, the graph is built during the optimization as proposed by .
Class name: GraphDiffusionRegulariser
AirLab includes a rich family of optimizers which are available in PyTorch including LBFGS, ASGD and Adam. They are tailored to optimize functions with a high number of parameters and thus are well suited for non-linear image registration objectives. We refer to  for a detailed overview of first order gradient based optimizers. As PyTorch also supports no-grad computations, iteration schemes as used in the Demons algorithm are also supported. The following snippet is an example usage of no-grad taken from the Demons regularizer.
Iii-C Upcoming Features
In this section, we list the features which did not make it into the present version, which however are planned for integration into AirLab until 30.9.2018.
Image domains: currently, AirLab supports only images which have the same domain with equal pixel spacing. This will be generalized to images with different domains and different pixel spacing.
Linear transforms: the similarity and affine transform will be integrated to also support scaling and shearing.
Interpolation: B-spline interpolation for the image warping is planned for integration.
In this section, we provide image registration examples. We have implemented two classic registration algorithms within AirLab and show their qualitative performance on synthetic examples and on a DirLab dataset . Quantitative analyses will follow in the final version of this paper.
Iv-a Image Registration Algorithms
The following algorithms have been implemented:
Rigid: a simple objective with a rigid transformation has been set up, where the similarity metric has been optimized with Adam.
FFD: the Free Form Deformations algorithm  was implemented. As in the original paper, a third order B-spline kernel has been used for the parametric transformation model. Furthermore, the similarity measure with the regularizer on the displacement field have been applied. The overall objective has been optimized with Adam.
Demons: the Demons algorithm  was implemented using the similarity measure with the Gaussian Demons regularizer. The similarity measure has been optimized with Adam, while the regularizer was applied after each iteration.
In FFD and Demons, a multi-resolution strategy has been implemented performing iterations for the FFD and iterations for the Demons algorithm. The detailed parameter configuration can be found in the source-code. The following snippet illustrates how to setup a registration algorithm in AirLab with the Rigid registration example:
Iv-A1 Rigid Example
For the Rigid example two AirLab images have been registered, where the moving image has been rotated. In Figure 1, the registration result is depicted.
Iv-A2 Demons Example
The Demons algorithm has been applied to the circle and C example. For better illustration, see Figure 2, a shaded circle has been warped with the final transformation.
Iv-A3 FFD Example
The same example has been registered also using the Demons algorithm (see Figure 4).
Iv-B Performance Analysis
All experiments have been conducted using an NVIDIA GeForce GTX 1080 GPU. We evaluate the performance of AirLab by profiling with the autograd profiler of PyTorch. One iteration on the last resolution level () of the DirLab examples have been executed. The used CPU and GPU time is listed in Table I.
Because the GaussianRegulariser is not differentiated, there is less computational time spent by autograd for the Demons example.
We have introduced AirLab, an environment for rapid prototyping and reproduction of medical image registration algorithms. It is written in the scripting language Python and heavily uses functionality of PyTorch. The unique feature compared to existing image registration software is the automatic differentiation which fosters rapid prototyping. AirLab is freely available under the Apache License 2.0 and accessible on GitHub: https://github.com/airlab-unibas/airlab.
With AirLab, we hope that we can make a valuable contribution to the medical image registration community, and we are looking forward to see researchers and developers which activly use AirLab in their work. Finally, we encourage them also to contribute to future development of AirLab.
The authors thank Alina Giger, Reinhard Wendler and Simon Pezold for their great support.
-  Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In OSDI, volume 16, pages 265–283, 2016.
-  Brian B Avants, Nicholas J Tustison, Gang Song, Philip A Cook, Arno Klein, and James C Gee. A reproducible evaluation of ants similarity metric performance in brain image registration. Neuroimage, 54(3):2033–2044, 2011.
-  Ghazi Bouabene, Thomas Gerig, Marcel Lüthi, Andreas Forster, Dennis Madsen, and Dana Rahbani. Scalismo - scalable image analysis and shape modelling. https://github.com/unibas-gravis/scalismo, 2015. [Online; accessed 25-June-2018].
-  Richard H Byrd, Jorge Nocedal, and Robert B Schnabel. Representations of quasi-newton matrices and their use in limited memory methods. Mathematical Programming, 63(1-3):129–156, 1994.
-  Pascal Cachier and Xavier Pennec. 3d non-rigid registration by gradient descent on a gaussian-windowed similarity measure using convolutions. In Mathematical Methods in Biomedical Image Analysis, 2000. Proceedings. IEEE Workshop on, pages 182–189. IEEE, 2000.
-  Richard Castillo, Edward Castillo, Rudy Guerra, Valen E Johnson, Travis McPhail, Amit K Garg, and Thomas Guerrero. A framework for evaluation of deformable image registration spatial accuracy using large landmark point sets. Physics in Medicine & Biology, 54(7):1849, 2009.
-  Vincent Dumoulin and Francesco Visin. A guide to convolution arithmetic for deep learning. arXiv preprint arXiv:1603.07285, 2016.
-  Ben Glocker, Nikos Komodakis, Georgios Tziritas, Nassir Navab, and Nikos Paragios. Dense image registration through mrfs and efficient linear programming. Medical image analysis, 12(6):731–741, 2008.
-  Eldad Haber and Jan Modersitzki. Intensity gradient based registration and fusion of multi-modal images. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 726–733. Springer, 2006.
-  Berthold KP Horn and Brian G Schunck. Determining optical flow. Artificial intelligence, 17(1-3):185–203, 1981.
-  Christoph Jud, Nadia Möri, Benedikt Bitterli, and Philippe C Cattin. Bilateral regularization in reproducing kernel hilbert spaces for discontinuity preserving image registration. In International Workshop on Machine Learning in Medical Imaging, pages 10–17. Springer, 2016.
-  Christoph Jud, Nadia Mori, and Philippe C Cattin. Sparse kernel machines for discontinuous registration and nonstationary regularization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pages 9–16, 2016.
-  Christoph Jud, Robin Sandkühler, and Philippe C Cattin. An inhomogeneous multi-resolution regularization concept for discontinuity preserving image registration. In International Workshop on Biomedical Image Registration, pages 3–12. Springer, 2018.
-  András P Keszei, Benjamin Berkels, and Thomas M Deserno. Survey of non-rigid registration tools in medicine. Journal of digital imaging, 30(1):102–116, 2017.
-  Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
-  Stefan Klein, Marius Staring, Keelin Murphy, Max A Viergever, and Josien PW Pluim. Elastix: a toolbox for intensity-based medical image registration. IEEE transactions on medical imaging, 29(1):196–205, 2010.
-  Julian Krebs, Tommaso Mansi, Boris Mailhé, Nicholas Ayache, and Hervé Delingette. Learning structured deformations using diffeomorphic registration. arXiv preprint arXiv:1804.07172, 2018.
-  Bradley Christopher Lowekamp, David T Chen, Luis Ibáñez, and Daniel Blezek. The design of simpleitk. Frontiers in neuroinformatics, 7:45, 2013.
-  Marcel Lüthi, Thomas Gerig, Christoph Jud, and Thomas Vetter. Gaussian process morphable models. IEEE transactions on pattern analysis and machine intelligence, 2017.
-  Frederik Maes, Andre Collignon, Dirk Vandermeulen, Guy Marchal, and Paul Suetens. Multimodality image registration by maximization of mutual information. IEEE transactions on Medical Imaging, 16(2):187–198, 1997.
-  Kasper Marstal, Floris Berendsen, Marius Staring, and Stefan Klein. Simpleelastix: A user-friendly, multi-lingual library for medical image registration. In Computer Vision and Pattern Recognition Workshops (CVPRW), 2016 IEEE Conference on, pages 574–582. IEEE, 2016.
-  Jan Modersitzki. FAIR: flexible algorithms for image registration, volume 6. SIAM, 2009.
-  Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W, 2017.
-  Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747, 2016.
-  Daniel Rueckert, Luke I Sonoda, Carmel Hayes, Derek LG Hill, Martin O Leach, and David J Hawkes. Nonrigid registration using free-form deformations: application to breast mr images. IEEE transactions on medical imaging, 18(8):712–721, 1999.
-  Robin Sandkühler, Christoph Jud, Simon Pezold, and Philippe C Cattin. Adaptive graph diffusion regularisation for discontinuity preserving image registration. In Biomedical Image Registration. Springer, 2018.
-  Julia A Schnabel, Daniel Rueckert, Marcel Quist, Jane M Blackall, Andy D Castellano-Smith, Thomas Hartkens, Graeme P Penney, Walter A Hall, Haiying Liu, Charles L Truwit, et al. A generic framework for non-rigid registration based on non-uniform multi-level free-form deformations. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 573–581. Springer, 2001.
-  J-P Thirion. Image matching as a diffusion process: an analogy with maxwell’s demons. Medical image analysis, 2(3):243–260, 1998.
-  Tom Vercauteren, Xavier Pennec, Aymeric Perchant, and Nicholas Ayache. Non-parametric diffeomorphic image registration with the demons algorithm. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 319–326. Springer, 2007.
-  M.A. Viergever, J.B.A. Maintz, S. Klein, K. Murphy, M. Staring, and J.P.W. Pluim. A survey of medical image registration - under review. Medical Image Analysi, 33:140–144, 10 2016.
-  Paul Viola and William M Wells III. Alignment by maximization of mutual information. International journal of computer vision, 24(2):137–154, 1997.
-  Valery Vishnevskiy, Tobias Gass, Gabor Szekely, Christine Tanner, and Orcun Goksel. Isotropic total variation regularization of displacements in parametric image registration. IEEE transactions on medical imaging, 36(2):385–395, 2017.
-  Roger P Woods, Scott T Grafton, Colin J Holmes, Simon R Cherry, and John C Mazziotta. Automated image registration: I. general methods and intrasubject, intramodality validation. Journal of computer assisted tomography, 22(1):139–152, 1998.
-  Terry S Yoo, Michael J Ackerman, William E Lorensen, Will Schroeder, Vikram Chalana, Stephen Aylward, Dimitris Metaxas, and Ross Whitaker. Engineering and algorithm design for an image processing api: a technical report on itk-the insight toolkit. Studies in health technology and informatics, pages 586–592, 2002.
We finally list the MSELoss similarity metric and the TVRegulariser class as snippet to illustrate how similarity measures and regularizers are implemented in AirLab. (Please see the following page)