[论文解读] AirLab: Autograd Image Registration Laboratory

Robin Sandkühler, Christoph Jud+2 more|arXiv (Cornell University)|Jun 26, 2018

Medical Image Segmentation Techniques参考文献 41被引用 49

一句话总结

AIRLab 提供一个基于 Python、由 PyTorch 驱动的环境，用于快速原型设计和医学图像配准的可重复性，具备自动微分和 GPU 加速执行。

ABSTRACT

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 presented draft of this paper outlines AIRLab with first code snippets and performance analyses. A more exhaustive introduction will follow as a final version soon.

研究动机与目标

• 推动在医学图像配准中对快速原型化环境的需求，以应对复杂目标和大规模的3D数据。
• 介绍 AIRLab 作为一个开放的、可微分的框架，使配准目标能够自动计算梯度。
• 提供构建模块（相似性度量、变换、正则化器）以及示例方法，以促进快速试验。
• 通过初步实验和代码示例演示性能和可重复性的优势。

提出的方法

• 通过 PyTorch autograd 对配准目标进行自动微分。
• 一组模块化的相似性度量（MSE、NCC、LCC、SSIM、MI、NGF）和正则化项（扩散、各向异性/等向性 TV、稀疏性、Demons 等）用于 f 及其参数。
• 支持线性/密集、非线性/插值（例如 B-spline 和 Wendland 核）以及密集/混合变换模型。
• 反向图像扭曲以在固定图像域中计算残差，结合 GPU 加速的网格采样。
• 可选的可微分同胚变换，使用带指数映射的静止速度场。
• 一个用于在如 POPI 之类数据集上自动评估配准的评估流水线。
• 可用的来自 PyTorch 的常用优化器（如 LBFGS、Adam），用于高参数问题。
• 针对不同图像尺寸的 CPU vs. GPU 执行性能基准。

实验结果

研究问题

• RQ1 自动微分是否可以简化配准目标的快速原型设计，并在无需手动梯度推导的情况下实现可重复的实验？
• RQ2 不同的相似性度量和正则化项如何影响 2D/3D 数据的配准性能和稳定性？
• RQ3 对于不同的变换模型，AIRLab 在 CPU 与 GPU 上的实际性能特征（速度、可扩展性）是多少？
• RQ4 在一个统一的可微分框架中，AIRLab 在多大程度上能够复现刚性、自由形变（FFD）和可微分 Demons 等标准配准方法？

主要发现

• AIRLab 使配准目标能够自动计算梯度，减少手动实现工作量。
• 它提供了多种相似性度量和正则化项，以及若干变换模型，便于快速原型化配准方法。
• 初步实验展示了在 AIRLab 内对刚性、FFD 和 Diffeomorphic Demons 的配准工作流及相应的可视化结果。
• 性能统计显示，对于较大的 2D 输入，基于 GPU 的运行大幅减少计算时间，并展示了随图像尺寸的扩展性。
• 包含一个自动评估流水线，以促进在如 POPI 这类数据集上的可重复性。
• 本文给出可执行的代码片段，并计划整合如 B-spline 插值和多通道数据等额外特征。