[論文レビュー] TorchRadon: Fast Differentiable Routines for Computed Tomography

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

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

ひとこと要約

TorchRadonはCUDAライブラリで、微分可能なRadon変換、shearlets、そしてCTソルバーをPyTorchで提供し、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再構成を組み合わせる動機づけ。
• PyTorchのトレーニンググラフに統合可能な微分可能なCT演算子を提供。
• CTルーチンと反復アルゴリズムを加速する高速なGPUバックエンドを提供。

提案手法

• 微分可能な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のバッチ処理とメモリ効率的なテクスチャにより、Astra Toolboxより前方/後向のRadon投影を最大125x速く実現。
• バッチ処理とGPUベースのカーネルにより、TorchRadonはCPU/GPUベースのベースラインを大きく上回る速度を達成（例：いくつかのバッチシナリオで40x超のスピードアップ）。
• 半精度ストレージはRadon変換とFBP再構成で精度低下が限定的なまま、 substantialな速度向上（2.5x超）をもたらす。
• 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演算子を実現します。