摘要

去年，我们尝试使用 Dask/ITK/Scikit-Image 对一组 3D 图像进行大规模图像分析。具体来说，我们研究了反卷积，这是一种常见的图像去模糊方法。一年后的现在，我们带着对 Dask 和 CuPy 如何交互的更好理解、增强的序列化方法以及开源社区的支持，再次回顾这些实验。本文将介绍以下内容

1. 实现用于 CPU + GPU 的通用反卷积方法
2. 利用 Dask 在更大的数据集上执行反卷积
3. 使用 Napari 图像查看器探索结果

图像分析重温

之前我们使用了来自 ITK 和 Scikit-Image 的 Richardson Lucy (RL) 反卷积算法。我们停留在理论上，探讨 GPU 如何潜在地帮助加速这些工作流。从 Scikit-Image 的实现开始，我们天真地尝试用 cupyx.scipy.ndimage.convolve 替换 scipy.signal.convolve 调用，虽然性能有所提升，但并没有显著提升——也就是说，我们没有达到想要的 100 倍加速。

在数学中，通常一个在某种表示下求解效率低下的问题，可以通过预先对数据进行转换来提高效率。在这种新的表示下，我们可以在将结果转换回更熟悉的表示之前，更容易地解决相同的问题（在本例中是卷积）。说到卷积，我们应用的转换称为 快速傅里叶变换 (FFT)。应用此转换后，我们可以使用简单的乘法对数据进行卷积。

事实证明，FFT 变换在 CPU 和 GPU 上都极快。同样，我们可以用 FFT 编写的算法也会加速。这是图像处理领域常用的技术，用于加速卷积。尽管增加了进行 FFT 的步骤，但转换成本 + 算法成本仍然低于在原始空间中执行算法。我们（以及之前的其他人）发现 Richardson Lucy（在 CPU 和 GPU 上）就是这种情况，并且当我们使用 Dask 在多个 GPU 上并行化时，性能持续提升。

来自开源社区的帮助

FFT RL 等价算法已经存在了一段时间，太阳动力学观测站的好心人构建并分享了一个 NumPy/CuPy 实现，作为 大气成像组件 Python 包 (aiapy) 的一部分。我们稍微修改了他们的实现，使其能够处理 3D 和 2D 点扩散函数，并利用 NEP-18 方便地将 NumPy 和 CuPy 调度到 NumPy 和 CuPy 函数。

def deconvolve(img, psf=None, iterations=20)
# 用零填充 PSF，使其形状与图像匹配
pad_l, pad_r = np.divmod(np.array(img.shape) - np.array(psf.shape), 2)
pad_r += pad_l
psf = np.pad(psf, tuple(zip(pad_l, pad_r)), 'constant', constant_values=0)

# 将 PSF 中心移到原点
# 需要确保 PSF 在进行卷积时不会引入偏移
# 与图像进行卷积
for i in range(psf.ndim)
psf = np.roll(psf, psf.shape[i] // 2, axis=i)

# 卷积需要 PSF 的 FFT
psf = np.fft.rfftn(psf)

# 在图像副本上进行就地反卷积
# (避免改变原始图像)
img_decon = np.copy(img)
for _ in range(iterations)
ratio = img / np.fft.irfftn(np.fft.rfftn(img_decon) * psf)
img_decon *= np.fft.irfftn((np.fft.rfftn(ratio).conj() * psf).conj())
return img_decon

对于 1.3 GB 的图像，我们测量了以下结果

• CuPy：20 次迭代约 3 秒
• NumPy：2 次迭代约 36 秒

我们看到，在迭代次数增加 10 倍的情况下，速度提高了 10 倍——非常接近我们期望的 100 倍加速！让我们来看看这个实现如何处理真实的生物数据和 Dask…

定义 Dask 集群并加载数据

我们获得了 NIH Shroff 教授实验室提供的样本数据。原始数据是 3D TIFF 文件，我们随后将其转换为形状为 (950, 2048, 2048) 的 Zarr 格式。

我们首先在 DGX2（单个节点包含 16 个 GPU）上创建一个 Dask 集群，并加载存储在 Zarr 中的图像

示例 Notebook

from dask.distributed import Client
from dask_cuda import LocalCUDACluster
import dask.array as da

import rmm
import cupy as cp

cluster = LocalCUDACluster(
local_directory="/tmp/bzaitlen",
enable_nvlink=True,
rmm_pool_size="26GB",
)
client = Client(cluster)

client.run(
cp.cuda.set_allocator,
rmm.rmm_cupy_allocator
)

imgs = da.from_zarr("/public/NVMICROSCOPY/y1z1_C1_A.zarr/")

Array Chunk Bytes 7.97 GB 8.39 MB Shape (950, 2048, 2048) (1, 2048, 2048) Count 951 Tasks 950 Chunks Type uint16 numpy.ndarray 20482048950

从上面的 Dask 输出可以看出，数据是包含 950 个图像的 z 堆栈，每个切片都是 2048x2048。对于这个数据集，如果我们使用更大的块进行操作，可以提高 GPU 性能。此外，我们需要确保块的大小至少与 PSF（在本例中为 (128, 128, 128)）一样大。由于我们在拥有 16 个 GPU 的 DGX2 上进行工作，因此我们可以根据需要重新分块，轻松地将数据放入内存并在每个 GPU 上执行反卷积。

# 根据 PSF 形状分块 (128, 128, 128)
imgs = imgs.rechunk(chunks={0: 190, 1: 512, 2: 512})
imgs

Array Chunk Bytes 7.97 GB 99.61 MB Shape (950, 2048, 2048) (190, 512, 512) Count 967 Tasks 80 Chunks Type uint16 numpy.ndarray

接下来，我们将数据转换为 float32，因为数据可能不是浮点类型。此外，计算时 32 位比 64 位稍快，并且节省内存。下面我们将 cupy.asarray 映射到每个数据块上。cupy.asarray 将数据从主机内存 (NumPy) 移动到设备/GPU (CuPy)。

imgs = imgs.astype(np.float32)
c_imgs = imgs.map_blocks(cp.asarray)

Array Chunk Bytes 15.94 GB 199.23 MB Shape (950, 2048, 2048) (190, 512, 512) Count 80 Tasks 80 Chunks Type float32 cupy.ndarray 20482048950

现在我们得到的是一个由 16 个 CuPy 数据块组成的 Dask 数组。请注意 Dask 在 SVG 输出中提供了很好的类型信息。当我们从 NumPy 切换到 CuPy 时，上面的框图显示类型为：cupy.ndarray——这是一个很好的健全性检查。

在运行反卷积之前，我们需要的最后一个组件是 PSF，它也应该加载到 GPU 上

import skimage.io

psf = skimage.io.imread("/public/NVMICROSCOPY/PSF.tif")
c_psf = cp.asarray(psf)

最后，我们在 Dask 数组上使用 deconvolve 函数调用 map_overlap

out = da.map_overlap(
deconvolve,
c_imgs,
psf=c_psf,
iterations=100,
meta=c_imgs._meta,
depth=tuple(np.array(c_psf.shape) // 2),
boundary="periodic"
)
out

上图取自小鼠肠道。

使用 Dask 和多个 GPU，我们测量了对 16GB 图像进行反卷积的时间约为 30 秒！但这仅仅是加速图像科学的第一步。

Napari

反卷积仅仅是图像科学家或显微镜学家需要的一个操作和工具。在研究潜在生物学时，他们还需要其他工具。在进入这些下一步之前，他们需要工具来可视化数据。Napari 是 PyData Bio 生态系统中使用的多维图像查看器，是可视化这些数据的好工具。作为一项实验，我们在连接了 NVLink 的本地工作站（配有 2 个 Quadro RTX 8000 GPU）上运行了相同的工作流。示例 Notebook

通过在我们的数组中添加 map_blocks 调用，我们可以将数据从 GPU 移回 CPU（从设备到主机）。

def cupy_to_numpy(x)
import cupy as cp
return cp.asnumpy(x)

np_out = out.map_blocks(cupy_to_numpy, meta=out)

当用户移动 Napari UI 上的滑块时，我们指示 dask 执行以下操作

• 从磁盘加载数据到 GPU (CuPy)
• 计算反卷积
• 移回主机 (NumPy)
• 使用 Napari 渲染

这大约有 1 秒的延迟，对于一个朴素的实现来说已经很不错了！我们可以通过添加缓存、改善与 map_overlap 的通信以及优化反卷积内核来改进这一点。

结论

现在我们展示了如何使用 Dask + CuPy 进行 Richardson-Lucy 反卷积。这只需要极少量的代码。结合图像查看器 (Napari)，我们能够检查数据和结果。所有这些都通过整合 PyData 库（Dask、CuPy、Zarr 和 Napari）以及一个新的反卷积内核而表现良好。希望这能为您开始分析自己的数据提供一个良好的模板，并展示自定义工作流的丰富性和易于表达性。如果您遇到任何挑战，请随时在 Dask issue tracker 上与我们联系，我们将很乐意与您交流 :)