论文标题
在TPU上加速MRI重建
Accelerating MRI Reconstruction on TPUs
论文作者
论文摘要
高级磁共振(MR)图像重建(例如压缩传感和基于子空间的成像)被认为是大规模,迭代,优化问题。鉴于实际临床使用所需的大量重建,这些高级重建方法的计算时间通常是不可接受的。在这项工作中,我们建议使用Google的张量处理单元(TPU)加速MR图像重建。 TPU是用于机器学习应用的特定应用集成电路(ASIC),最近已用于解决大规模的科学计算问题。作为概念验证,我们在Tensorflow中实现了乘数的交替方向方法(ADMM),以重建TPU上的图像。重建基于多通道,稀疏采样,以及带有稀疏约束的径向 - trajectory $ k $ - 空间数据。正向和非均匀傅立叶变换操作是根据离散傅立叶变换中的矩阵乘法制定的。稀疏的转换及其伴随作业被提出为卷积。数据分解应用于测量的$ K $空间数据,以便上述张量操作位于单个TPU内核中。数据分解和核心间通信策略是根据TPU互连网络拓扑设计的,以最大程度地减少通信时间。通过数值示例证明了所提出的基于TPU的图像重建方法的准确性和高平行效率。
The advanced magnetic resonance (MR) image reconstructions such as the compressed sensing and subspace-based imaging are considered as large-scale, iterative, optimization problems. Given the large number of reconstructions required by the practical clinical usage, the computation time of these advanced reconstruction methods is often unacceptable. In this work, we propose using Google's Tensor Processing Units (TPUs) to accelerate the MR image reconstruction. TPU is an application-specific integrated circuit (ASIC) for machine learning applications, which has recently been used to solve large-scale scientific computing problems. As proof-of-concept, we implement the alternating direction method of multipliers (ADMM) in TensorFlow to reconstruct images on TPUs. The reconstruction is based on multi-channel, sparsely sampled, and radial-trajectory $k$-space data with sparsity constraints. The forward and inverse non-uniform Fourier transform operations are formulated in terms of matrix multiplications as in the discrete Fourier transform. The sparsifying transform and its adjoint operations are formulated as convolutions. The data decomposition is applied to the measured $k$-space data such that the aforementioned tensor operations are localized within individual TPU cores. The data decomposition and the inter-core communication strategy are designed in accordance with the TPU interconnect network topology in order to minimize the communication time. The accuracy and the high parallel efficiency of the proposed TPU-based image reconstruction method are demonstrated through numerical examples.