torchkbnufft

文档 | GitHub | 笔记本示例

从PyPI简单安装

pip install torchkbnufft

关于

torchkbnufft 在PyTorch中实现了具有Kaiser-Bessel网格的非均匀快速傅里叶变换 [1, 2]。实现完全用Python编写，无需编译，便于以可读的代码灵活部署。NUFFT函数每个都封装为 torch.autograd.Function，允许通过NUFFT算子进行反向传播以训练神经网络。

此包在很大程度上受到了密歇根图像重建工具箱 (Matlab) 中NUFFT实现的启发。

操作模式和阶段

该包具有三种主要的NUFFT操作模式：基于表的NUFFT插值、基于稀疏矩阵的NUFFT插值以及具有Toeplitz嵌入FFT的前向/后向算子 [3]。大致来说，计算速度遵循以下顺序

通常最好从表插值开始，然后尝试其他模式以适应您的问题。

可以通过将它们传递到 KbNufft 或 KbNufftAdjoint 对象中来合并敏感性图。

文档

一个基于HTML的文档参考，可在Read the Docs上阅读。

大多数文件都附有文档字符串，可以在运行IPython时使用help阅读。示例

from torchkbnufft import KbNufft

help(KbNufft)

示例

torchkbnufft可用于N维NUFFT变换。这里的示例从一个简单的2D NUFFT开始，然后扩展到SENSE（一个具有多个并行2D NUFFT的任务）。

最后两个示例展示了基于稀疏矩阵乘法的NUFFT（对于高维情况可能很有用）和Toeplitz NUFFT（这是一种非常快速的正向反向NUFFT技术）。

所有示例都附有可以在Google Colab中运行的笔记本

• Colab中的基本示例
• Colab中的SENSE-NUFFT示例
• Colab中的稀疏矩阵示例
• Colab中的Toeplitz示例

简单的正向NUFFT

Colab中的基本NUFFT示例

以下代码加载了Shepp-Logan幻象并计算了单根径向k空间数据

import torch
import torchkbnufft as tkbn
import numpy as np
from skimage.data import shepp_logan_phantom

x = shepp_logan_phantom().astype(np.complex)
im_size = x.shape
# convert to tensor, unsqueeze batch and coil dimension
# output size: (1, 1, ny, nx)
x = torch.tensor(x).unsqueeze(0).unsqueeze(0).to(torch.complex64)

klength = 64
ktraj = np.stack(
    (np.zeros(64), np.linspace(-np.pi, np.pi, klength))
)
# convert to tensor, unsqueeze batch dimension
# output size: (2, klength)
ktraj = torch.tensor(ktraj, dtype=torch.float)

nufft_ob = tkbn.KbNufft(im_size=im_size)
# outputs a (1, 1, klength) vector of k-space data
kdata = nufft_ob(x, ktraj)

SENSE-NUFFT

Colab中的SENSE-NUFFT示例

该包还包括用于处理SENSE-NUFFT算子的实用工具。上面的代码可以修改以包含敏感性图。

smaps = torch.rand(1, 8, 400, 400) + 1j * torch.rand(1, 8, 400, 400)
sense_data = nufft_ob(x, ktraj, smaps=smaps.to(x))

此代码首先在smaps中对敏感性线圈进行乘法，然后计算每个线圈64长度的径向射线。所有操作都广播到线圈中，这最小化了与Python解释器的交互，有助于提高计算速度。

稀疏矩阵预计算

Colab中的稀疏矩阵示例

稀疏矩阵是表插值的替代品。它们的速度可能会有所不同，但比标准表模式略微准确。以下代码计算稀疏插值矩阵并使用它们来计算单根径向k空间数据

adjnufft_ob = tkbn.KbNufftAdjoint(im_size=im_size)

# precompute the sparse interpolation matrices
interp_mats = tkbn.calc_tensor_spmatrix(
    ktraj,
    im_size=im_size
)

# use sparse matrices in adjoint
image = adjnufft_ob(kdata, ktraj, interp_mats)

PyTorch只对实数实现了稀疏矩阵乘法，这可能会限制它们的速度。

Toeplitz嵌入

Colab中的Toeplitz示例

该包包括计算嵌入Toeplitz核的程序，并将它们用作正向/反向NUFFT操作的FFT滤波器[3]。这对于必须使用正向/反向操作来计算梯度的梯度下降算法非常有用。以下代码展示了示例

toep_ob = tkbn.ToepNufft()

# precompute the embedded Toeplitz FFT kernel
kernel = tkbn.calc_toeplitz_kernel(ktraj, im_size)

# use FFT kernel from embedded Toeplitz matrix
image = toep_ob(image, kernel)

在GPU上运行

此存储库中包含的所有示例都可以在GPU上运行，方法是在函数调用之前将NUFFT对象和数据发送到GPU，例如

adjnufft_ob = adjnufft_ob.to(torch.device('cuda'))
kdata = kdata.to(torch.device('cuda'))
ktraj = ktraj.to(torch.device('cuda'))

image = adjnufft_ob(kdata, ktraj)

如果所有对象的底层dtype和device不匹配，PyTorch将抛出错误。务必确保您的数据和NUFFT对象在正确的设备上并以正确的格式存在，以避免这些错误。

计算速度

以下秒数是在具有Xeon E5-2698 CPU和Nvidia Quadro GP100 GPU的工作站上观察到的计算时间。CPU计算限制在8个线程，并使用64位浮点数完成，而GPU计算使用32位浮点数完成。基准测试使用了torchkbnufft版本1.0.0和torch版本1.7.1。

(n) = 正常，(spm) = 稀疏矩阵，(toep) = Toeplitz嵌入，(f/b) = 正向/反向组合

通过运行

pip install -r dev-requirements.txt
python profile_torchkbnufft.py

其他包

对于对其他计算平台上的NUFFT实现感兴趣的用户，以下是一些其他项目的部分列表

1. TF KB-NUFFT (TensorFlow的KB-NUFFT)
2. SigPy (用于NumPy数组，Numba（CPU）和CuPy（GPU）后端)
3. FINUFFT (用于MATLAB，Python，Julia，C等，非常高效)
4. NFFT (用于Julia)
5. PyNUFFT (用于NumPy，还具有PyCUDA/PyOpenCL后端)

参考文献

1. 费舍尔，J. A.，& 萨顿，B. P. (2003). 使用最小-最大插值进行非均匀快速傅里叶变换. IEEE 信号处理汇刊，51(2)，560-574。

2. 贝蒂，P. J.，西岛，D. G.，& 保罗，J. M. (2005). 具有最小过采样比的快速网格重建. IEEE 医学图像汇刊，24(6)，799-808。

3. 费希丁格，H. G.，格，K.，& 斯托默，T. (1995). 非均匀采样理论中的有效数值方法. 数值数学，69(4)，423-440。

引用

如果您使用此包，请引用

@conference{muckley:20:tah,
  author = {M. J. Muckley and R. Stern and T. Murrell and F. Knoll},
  title = {{TorchKbNufft}: A High-Level, Hardware-Agnostic Non-Uniform Fast {Fourier} Transform},
  booktitle = {ISMRM Workshop on Data Sampling \& Image Reconstruction},
  year = 2020,
  note = {Source code available at https://github.com/mmuckley/torchkbnufft}.
}