在 Diamond Light Source 使用 NVIDIA Holoscan 加速印刷工作流程

Diamond Light Source 是英国知名的同步辐射加速器设施，它为科学家们提供了密集的 X 射线、红外线以及其他形式的光线，用于研究材料和生物结构。该设施设有 30 多个实验站或束线，并承载着一些世界上最先进、最复杂的科学研究项目。

I08-1 是 Diamond Light Source 的软 X 射线光束线，提供了一种称为分版成像的先进高分辨率成像技术，可提供纳米级分辨率的图像。分版成像使用一种计算成像方法，根据 X 射线束与样本相互作用产生的测量结果或衍射图案，以纳米级分辨率重建样本图像。

这对于许多生物结构(例如细胞中的线粒体和细胞器)以及材料科学样本中的内部结构和缺陷的纳米级特征成像至关重要。这种重建图像的过程非常强大，但可能会导致测量数据和查看图像之间存在显著差距。

I08-1 检测器每秒可处理 25 帧图像，而每秒可处理数千帧图像的检测器即将投入使用。这些传感器仪器需要在边缘进行加速计算。

更快的扫描速度能够研究更动态的工作流程。它们提高了实验的吞吐量，实时处理为用户提供实时反馈，以调整实验样本、检测器设置和探索样本以发现有趣的科学结果。

本文讨论了我们使用 I08-1 通过重构数据分析工作流来加快波束线实验数据的实时处理速度的工作。它还解决了一些关键挑战，例如目前以串行方式运行的通用分版印刷工作流，其中图像以 25 Hz 的帧率写入磁盘。

扫描完成后，启动实时处理流程，处理应用程序（PtyPy）可以处理完整的数据集。PtyPy 应用程序已针对 GPU 加速进行了优化，但 I/O 通信仍是实现更高吞吐量的主要瓶颈。

为了加速分版印刷工作流程，我们推出了 NVIDIA Holoscan，这是一个专为传感器处理设计的 SDK。它让科学家、研究人员和开发者能够更轻松地优化和扩展他们的传感器处理工作流程（如图 2 所示）。例如，在 Holoscan Operator 中使用 JAX 库来加速图像预处理。

借助 Holoscan，研究人员和开发者能够开发高性能、低延迟的传感器处理应用程序，这些应用程序可以使用熟悉语言的参考示例更轻松地进行扩展。

I08-1 的软 X 射线分版摄影仪使用 sCMOS 摄像头为分版摄影实验收集衍射数据。原始数据以形状帧的形式提供(2048, 2048)和类型uint16.在将这些数据输入迭代分版求解器应用 (PtyPy) 之前，我们会在每一原始帧中执行以下常规预处理任务：

• 背景暗流图像的减法
• 围绕中心裁剪并重新融合以减少重建时间

第一项任务需要获得清晰的衍射图像，而另外两项任务则根据实验情况提供无损压缩水平。理想情况下，所有这些步骤都尽可能在靠近源的地方执行(例如，片上或使用 FPGA)。遗憾的是，在此场景中使用的特定摄像头都没有这些选项。

JAX 用于显著加快单线程 Python 脚本执行上述任务的速度，并尽可能减少对代码的更改。由于原始帧处理代码是用 NumPy 编写的，因此 JAX JIT 能够将处理例程融合到单个 GPU 内核中，与原始 NumPy 版本相比，单张图像的速度提高了 2000 倍以上（忽略了从主机到设备所需的数据传输）。即使考虑到数据传输，速度也比原始基于 CPU 的 NumPy 实现提高了 40 倍以上。

虽然获取数据的速度相对较快，但将数据重建为研究人员或研究人员可以解读的图像可能很容易需要几分钟或几十分钟。在显示扫描结果之前，扫描和图像之间的这种无操作时间效率很低，会影响研究人员确定仪器设置是否正确，或者所扫描的样本区域是否有趣的能力。

通过将 ptychography 应用程序构建为应用程序和 Holoscan 运算符片段的集合，开发者可以利用 PtyPy ptychography 代码、Holoscan AI 推理和 Holoscan 网络运算符，相对更快地对实时处理 ptychography 应用程序的新 GPU 加速版本进行原型设计和测试。

挑战的下一部分是如何加快波束线 I-08 拼接成像工作流程的实时处理帧率，以应对当前和未来的 sCMOS 帧率。通过重叠串行工作流程步骤并使用 Holoscan，此波束线应该能够提供与传感器的帧率相匹配的实时处理。这将使波束线用户能够实时观察以拼接方式重建的样本图像。

3D 重建会产生一个更大的问题，对实时处理具有相同的要求。扩展多 GPU 和多节点处理以提供许多扫描的重叠并行处理，可能是满足实时处理的处理和 I/O 要求的一种方法。

我们的合作旨在使用两个 NVIDIA A2 GPU 在本地边缘服务器上测试各种工作流配置，其中预处理在一个 GPU 上运行，图像重建在第二个 GPU 上运行。这种方法能够专注于定制的分版印刷代码，同时利用边缘网络 I/O 运算符和 AI 加速库，这些库可以轻松重复使用，并在必要时扩展到多节点以用于生产用途。

Holoscan 支持创建端到端数据串流工作流，可在 I08-1 光束线上进行实时印刷图像处理，从而显著丰富整体用户交互。 如前所述，其他 Diamond 光源波束线在千赫兹检测范围内运行，但没有一条能够以该速率执行实时处理。

在扫描数据上训练 AI 模型，然后使用模型在 GPU 上以波束线运行推理，是一种有望以千赫兹速度实现实时处理的方法。

总结

传感器处理流程(如本文中描述的分版印刷流程)将重要的处理和 I/O 要求整合到单个应用程序中。随着传感器分辨率和刷新率的提高，基于文件的方法不再可行，这促使处理重新设计使用实时流式传输工作流。

这需要适当考虑端到端性能，这立即凸显出整个流程中的 I/O 瓶颈。 NVIDIA GPU 边缘系统结合使用 JAX、CuPy 和 CUDA 加速预处理和重建操作，从而获得必要的性能。但是，根据我们的端到端分析，这只会放大 I/O 瓶颈的影响。

Holoscan 提供了用于构建流处理软件流程的工具，这些工具还可以利用硬件的功能。这包括操作人员直接从网络将数据提取到 GPU 中(反向)，以更好地输入 GPU 并提高其利用率。在纯流处理(如前面突出显示的预处理步骤)期间，这种利用率可能特别低。由于 Holoscan 的设计易于模块化，因此它还支持其他功能，例如所讨论的深度学习或实时可视化功能。

通过将 GPU 加速计算与 Holoscan 结合使用，I08-1 可以显著缩短处理 X 射线显微镜数据所需的时间，并加速图像处理的帧率。边缘节点配备了高性能计算 (HPC)硬件，包括安装在 CPU 服务器上的 NVIDIA GPU。这些边缘设备或服务器旨在加速图像处理和机器学习。

为实现实时处理，Diamond Light Source 采用分布式计算架构，其中包括多个边缘节点和一个中央数据处理设施。

边缘节点位于 X 射线束线附近，负责处理生成的数据。然后，处理后的数据将发送到中央数据处理设施，在那里进行进一步分析和存储。

基于 NVIDIA GPU 进行 AI 处理的边缘 HPC 可以多种方式加速 X 射线显微镜数据处理。GPU 在并行处理大量图像数据方面非常高效。在 X 射线显微镜检查中， NVIDIA GPU 可用于加速降噪、图像配准和图像分割等任务。

机器学习算法(例如深度学习神经网络)可用于分析 X 射线显微镜数据并提取有意义的信息。 NVIDIA GPU 非常适合加速这些算法的训练和推理阶段。

通过在流式 AI 框架中使用 NVIDIA Holoscan 进行图像处理、机器学习、断层成像重建和数据压缩，边缘 AI 处理可以显著减少处理 X 射线显微镜数据所需的时间和资源，并加快实现科学突破。

借助 HPC 边缘处理，Diamond Light Source 为科学家提供做出实时决策和加速研究所需的工具，从而朝着科学民主化迈出了重要的一步。