[论文解读] TorchRadon: Fast Differentiable Routines for Computed Tomography

Matteo Ronchetti|arXiv (Cornell University)|Sep 29, 2020

Medical Imaging Techniques and Applications参考文献 20被引用 39

一句话总结

TorchRadon 是一个 CUDA 库，提供在 PyTorch 中可微分的 Radon 变换、Shearlets 和 CT 求解器，在 Astra Toolbox 上实现高达125x的加速，并使其能够无缝集成到神经网络中。

ABSTRACT

This work presents TorchRadon -- an open source CUDA library which contains a set of differentiable routines for solving computed tomography (CT) reconstruction problems. The library is designed to help researchers working on CT problems to combine deep learning and model-based approaches. The package is developed as a PyTorch extension and can be seamlessly integrated into existing deep learning training code. Compared to the existing Astra Toolbox, TorchRadon is up to 125 faster. The operators implemented by TorchRadon allow the computation of gradients using PyTorch backward(), and can therefore be easily inserted inside existing neural networks architectures. Because of its speed and GPU support, TorchRadon can also be effectively used as a fast backend for the implementation of iterative algorithms. This paper presents the main functionalities of the library, compares results with existing libraries and provides examples of usage.

研究动机与目标

• 推动将深度学习与基于模型的 CT 重建相结合的动机。
• 提供可微分的 CT 操作符，能够集成到 PyTorch 的训练图中。
• 提供一个快速的 GPU 支撑后端，以加速 CT 例程和迭代算法。

提出的方法

• 将可微分的 Radon 前向/反向投影和 shearlet 变换实现为与 PyTorch autograd 兼容的算子。
• 提供用于投影（Radon、RadonFanbeam）和基于 GPU 的 Alpha shearlets 的子模块，并支持批处理。
• 包含可在神经网络中使用的迭代求解器（CG、CGNE、Landweber）以及正则化框架（基于 shearlet 的）。
• 与 Astra Toolbox 和 AlphaTransforms 进行基准测试，以展示速度和准确性的优势。
• 展示使用 GPU 纹理的半精度存储，以在保持精度的同时加速受内存瓶颈影响的 Radon 运算。

实验结果

研究问题

• RQ1 TorchRadon 构建的可微分 CT 运算符在速度和精度上与现有库（如 Astra Toolbox 和 AlphaTransforms）相比如何？
• RQ2 Radon 投影、shearlet 变换和 CT 求解器是否可以在不修改代码的情况下被有效地批处理并集成到 PyTorch 模型中？
• RQ3 在 GPU 设置中，半精度存储对 CT 重建的精度和速度有何影响？
• RQ4 TorchRadon 是否能够加速混合深度学习/基于模型的 CT 重建工作流（如 ADMM、展开网络）并实现端到端可微分？

主要发现

• TorchRadon 由于在 GPU 上的批处理和内存高效纹理，前向/后向 Radon 投影比 Astra Toolbox 快速达到最高 125x。
• 批处理和基于 GPU 的内核使 TorchRadon 在某些批处理场景中以远超 CPU/GPU 基线的速度提升（例如超过 40x 的加速）。
• 半精度存储带来显著的加速（超过 2.5x），Radon 变换和 FBP 重建的精度损失可以忽略。
• TorchRadon 的 shearlet 实现与 AlphaTransforms 相匹配，相对系数误差极小（在单精度下低至 3.86e-7），重建结果几乎完全相同（相对误差约 6e-16）。
• FBP 与迭代求解器（CGNE、Landweber）在 MSE/精度方面与基线相比具有竞争力甚至更优；例如，FBP 的 MSE 2.22e-4 与 Astra 的 2.02e-4 相比，CGNE 在一个基准中达到 4.42e-5。
• 该库与 PyTorch Autograd 无缝集成，在不修改训练代码的情况下，在神经网络模型中实现可微分的 CT 运算符。