2769

TensorFlow MRI: A Library for Modern Computational MRI on Heterogenous Systems
Javier Montalt-Tordera1, Jennifer Steeden1, and Vivek Muthurangu1
1University College London, London, United Kingdom

Synopsis

We present TensorFlow MRI, a new open-source library of TensorFlow operators for MR image reconstruction and processing. Its goal is to enable fast prototyping of modern MRI applications within a single computing framework. It is intended for researchers working in MR image reconstruction and/or processing, especially those interested in ML applications. The library is primarily Python-based and is easy to use, understand and extend. It has a single-command installation procedure and extensive documentation. Thanks to the use of a TensorFlow backend, it has excellent performance and runs on heterogenous and distributed systems.

Introduction

In recent years, machine learning (ML) has revolutionized research in MR image reconstruction and processing 1. One key factor behind this has been the development of free, open-source machine learning libraries, such as TensorFlow and PyTorch, which have made efficient neural network training and inference readily available to researchers.
Many modern MRI reconstruction and processing pipelines involve the use of neural networks, typically implemented with TensorFlow or PyTorch, but also need MRI-specific tools, for which users often rely on libraries such as Gadgetron 2, BART 3 and SigPy 4.
It is necessary to bridge the gap between ML frameworks and MRI reconstruction tools. To this end, BART recently added the ability to use pre-trained TensorFlow graphs and has natively implemented some neural networks 5. SigPy enables conversion of arrays and linear operators to their PyTorch equivalents. Gadgetron Inline AI enables ML inference within the Gadgetron streaming reconstruction 6. While these interfaces are very useful, they cannot provide the intimate link of a unified computing framework.

TensorFlow MRI

We present TensorFlow MRI, a library of TensorFlow operators for computational MRI which includes:

  • A fast, native non-uniform fast Fourier transform (see TensorFlow NUFFT).
  • A unified gateway for MR image reconstruction, which supports parallel imaging, compressed sensing, machine learning and partial Fourier methods.
  • Common linear and nonlinear operators, such as Fourier operators and regularizers, to aid in the development of novel image reconstruction techniques.
  • Multicoil imaging operators, such as coil combination, coil compression and estimation of coil sensitivity maps.
  • Calculation of non-Cartesian k-space trajectories and sampling density estimation.
  • A collection of Keras objects including models, layers, metrics, loss functions and callbacks for rapid development of neural networks.
  • Many other differentiable operators for common tasks such as array manipulation and image/signal processing.
The library has a Python interface and is mostly written in Python. However, computations are efficiently performed by the TensorFlow backend (implemented in C++/CUDA), which brings together the ease of use and fast prototyping of Python with the speed and efficiency of optimized lower-level implementations.
Being an extension of TensorFlow, TensorFlow MRI integrates seamlessly in ML applications. No additional interfacing is needed to include a SENSE operator within a neural network, or to use a trained prior as part of an iterative reconstruction. Therefore, the gap between ML and non-ML components of image processing pipelines is eliminated.
Whether an application involves ML or not, TensorFlow MRI operators can take full advantage of the TensorFlow framework, with capabilities including automatic differentiation, multi-device support (CPUs and GPUs), automatic device placement and copying of tensor data, and conversion to fast, serializable graphs.

TensorFlow NUFFT

The NUFFT is a basic operation in non-Cartesian MRI reconstruction. Iterative methods and neural networks operating with non-Cartesian k-space data will use it repeatedly. Therefore, an efficient NUFFT operator is essential.
Here we also present TensorFlow NUFFT. Implemented as a native op, it delivers a performance unlikely to be achievable with a composition of existing TensorFlow primitives. Internally we rely on the recent algorithms by the Flatiron Institute 7,8. TensorFlow NUFFT adapts this code, registers the CPU/GPU kernels to the TensorFlow framework, interfaces with the op kernel context and defines the gradient and shape inference functions. All of this is completely transparent to the user, who can use the NUFFT operator from Python just like any other TensorFlow operator.
Due to its standalone utility, TensorFlow NUFFT is released in its own package (see Practical Resources).

Examples and Performance

We evaluated the performance of the TensorFlow NUFFT operator (forward and adjoint) and compared it to BART and TF KB-NUFFT 9, another NUFFT operator for TensorFlow. Details and results are shown in Figure 1.
We also evaluated the performance of an iterative conjugate gradient SENSE reconstruction 10, implemented in TensorFlow MRI, and compared it with BART. Details and results are shown in Figure 2.
All measurements were performed on a workstation with an Intel Core i9-10980XE and 2x NVIDIA RTX 3090. Reported times are the average of 5 measurements, with 2 warm-up iterations.
Figure 3 shows an example of an iGRASP 11 reconstruction performed using only TensorFlow MRI. TensorFlow MRI has also been successfully deployed in online streaming reconstructions at our clinical scanners 12.

Practical Resources

TensorFlow MRI/NUFFT can be easily downloaded and installed using “pip”.
$ pip install tensorflow-mri
$ pip install tensorflow-nufft
TensorFlow NUFFT will be installed automatically for users of TensorFlow MRI, as it is one of its dependencies.
The TensorFlow MRI source code can be found on https://github.com/mrphys/tensorflow-mri and its documentation can be found on https://mrphys.github.io/tensorflow-mri. The TensorFlow NUFFT source code is on https://github.com/mrphys/tensorflow-nufft and its documentation is on https://mrphys.github.io/tensorflow-nufft.
Both packages are released under the permissive Apache 2.0 license, under which everyone is free to use, distribute and modify the software for research and commercial purposes.
Interested users are warmly invited to send us feedback, get involved and submit any issues and contributions at our GitHub pages.

Conclusion

We presented TensorFlow MRI, an open-source library that extends the TensorFlow ecosystem with MRI functionality and provides a unified computing framework for modern MRI reconstruction and processing, with rapid development in Python, efficient CPU and GPU computation and seamless integration of ML and non-ML methods.

Acknowledgements

We gratefully acknowledge funding support from UK Research and Innovation (UKRI) and Heart Research UK.

References

1. Lundervold AS, Lundervold A. An overview of deep learning in medical imaging focusing on MRI. Z. Med. Phys. 2019;29:102–127 doi: 10.1016/J.ZEMEDI.2018.11.002.
2. Hansen MS, Sørensen TS. Gadgetron: An open source framework for medical image reconstruction. Magn. Reson. Med. 2013;69:1768–1776 doi: 10.1002/MRM.24389.
3. Uecker M, Holme C, Blumenthal M, et al. mrirecon/bart: version 0.7.00. 2021 doi: 10.5281/ZENODO.4570601.
4. Ong F, Lustig M. SigPy: A Python package for high performance iterative reconstruction. In: Proceedings of the ISMRM 27th Annual Meeting & Exhibition. Montreal; 2019.
5. Blumenthal M, Uecker M. Deep, Deep Learning with BART. In: Proc. Intl. Soc. Mag. Reson. Med. 29. ; 2021.
6. Xue H, Davies R, Hansen D, et al. Gadgetron Inline AI: Effective Model Inference on MR Scanner. In: Proc. Intl. Soc. Mag. Reson. Med. 27. ; 2019.
7. Barnett AH, Magland J, af Klinteberg L. A Parallel Nonuniform Fast Fourier Transform Library Based on an “Exponential of Semicircle" Kernel. SIAM J. Sci. Comput. 2019;41:C479–C504 doi: 10.1137/18M120885X.
8. Shih Y, Wright G, Anden J, Blaschke J, Barnett AH. cuFINUFFT: a load-balanced GPU library for general-purpose nonuniform FFTs. 2021 IEEE Int. Parallel Distrib. Process. Symp. Work. 2021:688–697 doi: 10.1109/IPDPSW52791.2021.00105.
9. Muckley M, Stern R, Murrell T, Knoll F. TorchKbNufft: A High-Level, Hardware-Agnostic Non-Uniform Fast Fourier Transform. In: ISMRM Workshop on Data Sampling & Image Reconstruction. ; 2020.
10. Pruessmann KP, Weiger M, Börnert P, Boesiger P. Advances in sensitivity encoding with arbitrary k-space trajectories. Magn. Reson. Med. 2001;46:638–651 doi: 10.1002/mrm.1241.
11. Feng L, Grimm R, Block KT, et al. Golden-angle radial sparse parallel MRI: Combination of compressed sensing, parallel imaging, and golden-angle radial sampling for fast and flexible dynamic volumetric MRI. Magn. Reson. Med. 2014;72:707–717 doi: 10.1002/MRM.24980.
12. Jaubert O, Montalt-Tordera J, Arridge S, Steeden J, Muthurangu V. FlowSeg: Joint deep artifact suppression and segmentation network for low latency beat-to-beat assessment of cardiac aortic flow, Submitted to ISMRM 2022.

Figures

Figure 1. Execution times of forward and adjoint NUFFT. The test data was a 384 x 384 image with 600 radial spokes (460,800 nonuniform points). Three methods are compared: BART, TF KB-NUFFT and TensorFlow NUFFT. BART does not expose a GPU NUFFT operator (only available for iterative inverse NUFFT), so it could not be tested. Note that the BART method would not be accessible from TensorFlow applications.

Figure 2. Execution times of the CG-SENSE reconstruction. Tested on a 12-coil 256x256 image with 96 radial spokes (24,576 nonuniform points). The TensorFlow MRI reconstruction was run in eager and graph execution modes. Graph execution enables additional optimization. The conversion to graph has an overhead of ~1 s, but this only requires shape and type information and can be done offline. The reasons for the relatively poor performance of the BART GPU reconstruction are unknown to us.

Figure 3. iGRASP reconstruction performed with TensorFlow MRI. The test data is a liver dynamic contrast enhancement dataset (384 x 384 matrix, 21 radial spokes per frame, 28 frames, 12 coils). The regularization factor is 0.001.

Proc. Intl. Soc. Mag. Reson. Med. 30 (2022)
2769
DOI: https://doi.org/10.58530/2022/2769