天文台和望远镜、激光器和X射线光源,以及其他高通量仪器产生图像数据的速度,已经超过了依赖CPU的处理流程所能支持的决策时效。计算瓶颈很少只出现在某一个运算环节上,而是出现在从传感器数据到最终决策的整条链路上:读取原始数据、匹配点扩散函数或探测器响应、扣除或归约、拟合、分类、再发出告警。

现代科研设施在一次巡天或一轮实验活动中,往往会积累起数PB级的多维数据。如果采用以CPU为主的实现方式,从仪器在几秒钟内产生的数据中得出具有科学价值的结果,可能需要耗费数小时、数月,甚至一年以上的时间。这种延迟使科学家和工程师更难从数据中提炼洞见、推进研究进展。真正缺失的是一条从传感器到决策、全链路原生运行在GPU上的处理路径,而不仅仅是把某一个环节加速。

NVIDIA cuPhoton是一款开源的NVIDIA CUDA-X工具包,提供了面向这条处理路径的GPU加速基础模块,应用场景覆盖从光谱与光学天文学,到时域激光/X射线分析等多个领域。通过让图像数据从传感器读取到分类全程保留在GPU上,cuPhoton将原本需要长时间等待的处理过程压缩到秒级,形成一个研究人员可以交互式运行的紧密循环,从而缩短从数据到洞见的时间,帮助科研设施跟上仪器产生数据的速度。

用cuPhoton将科学洞见的获取时间从数月缩短到数分钟

在天文学、X射线科学以及其他高通量领域,NVIDIA cuPhoton提供了用于加载、处理和分析海量多维数据集的工具。下面的例子展示了如何将多个cuPhoton模块组合成一条端到端的工作流。

美国国家科学基金会与能源部共建的维拉·鲁宾天文台(Vera C. Rubin Observatory),是对快速提取科学洞见有着强烈需求的现代科研设施的典型代表。每到天黑之后,该天文台的LSSTCam每39秒就会拍摄一张新的、关于南天的32亿像素曝光图像。实时处理(Prompt Processing)流水线会将每一帧与同一天区的参考模板进行比对。在大约60到120秒内,系统必须判断约1万个探测目标是属于天体瞬变事件,还是宇宙射线、卫星轨迹、处理错误等伪影。整整一夜下来,这意味着多达20TB的图像数据和1000万个候选目标需要处理。

NVIDIA cuPhoton能够在多GPU、多节点的NVIDIA Grace Blackwell以及NVIDIA Vera Rubin系统上进行扩展。在使用多个GPU处理数百TB规模的代表性工作负载时,cuPhoton将图像加载与读取速度提升了最高14900倍,信号处理速度提升了最高14550倍,把原本需要数月的计算时间压缩到了几分钟。在GPU加速的Python环境下,原本需要九个月才能完成的数据分析,已经被证明可以在四小时内完成。对于规模较小、仅为千字节量级的数据集,整个cuPhoton处理流程可以在毫秒甚至微秒级的时间尺度上执行完毕。相关图表展示了cuPhoton各单项操作相对x86 CPU基准的加速效果;需要说明的是,这些数字反映的是单项操作的加速比,并非端到端流水线的整体加速比。

一个示例科学工作流的详细演示

本文以一组经过物理建模、具有真实感的合成天文图像为例,演示一个具有代表性的科学工作流。以下示例的规模经过精心设计,可以在NVIDIA DGX Spark或单GPU工作站上运行。同样的工作流也可以在多GPU工作站和集群上扩展,相关的启动方式将在后文给出。整个工作流大致包括以下几个阶段:

首先,将当前曝光图像和参考数据加载到GPU显存中。

对齐:将两次观测结果放置到同一套天球网格上,使同一颗星落在相同的像素位置。

匹配点扩散函数并做差:匹配当前曝光图像与同一天区更早参考图像的点扩散函数(PSF,即描述光在像素间如何分布的模式),然后将参考图像从当前曝光中扣除。一个运动中的天体会在差异图像上留下一对正负像素配对的图案,称为“偶极子”。

对该偶极子拟合双PSF模型,以估算运动天体的位置。

可视化结果:展示PSF匹配核函数、相减后的残差图像,以及还原出的偶极子图案。

工作流的具体执行过程

下面的示例使用小规模合成数据集,方便读者在自己的机器上进行测试。前文给出的峰值性能数字会因工作负载、数据集、具体实现方式和硬件配置的不同而变化。其中最大的加速倍数均来自单项操作,不应被理解为端到端流水线的整体加速系数。

本次演示涵盖cuPhoton的五个核心组件:xDataReader、xRep、xPois、xFit和xScan。该工具包还包含用于时域X射线探测器分析的xRay模块。

首先,克隆代码仓库,并按照README中的说明进行环境搭建。安装命令(uv sync --locked --extra dev --extra gpu --extra viz)会配置一套锁定版本的Linux环境,底层基于NVIDIA CUDA 13生态系统,包括CuPy、PyTorch、Numba-CUDA、KvikIO和NVIDIA nvCOMP,并附带Bokeh和Pillow等可视化库。同一份代码还包含cuPhoton命令行工具、本文用到的各个模块,以及examples/run_quickstarts.py脚本,该脚本会生成合成输入数据,并将运行产物写入quickstart-output/目录。

搭建工作目录

按照README文档创建锁定版本的CUDA 13 GPU环境。xDataReader的FITS处理路径还需要其原生扩展模块;在代码仓库根目录下执行bash src/cuphoton/xdr/src/build.sh(参见xDataReader指南)。cuPhoton 0.1.3版本支持Python 3.12至3.14,以及Linux上的CUDA 13。

从源码构建xDataReader需要C++17编译器、CUDA与cuFile开发头文件,以及支持可重入特性的CFITSIO开发库。具体安装步骤请参考xDataReader指南。

(具体命令略,参见原始仓库说明)

在同一个交互式会话(如Jupyter、IPython或集成脚本)中,从代码根目录开始,依次顺序执行每一段Python代码片段。后续章节会复用会话中保留的变量,包括images、aligned、reference、target、fit、stamp、model和result。请确保WORK_DIR在整个执行过程中始终处于已定义状态。如果原生xDataReader组件尚未构建,请先暂停并执行bash src/cuphoton/xdr/src/build.sh。

完整的工作流脚本可在cuPhoton代码仓库中找到。若要在全新终端中进行自动化运行,可通过uv run python调用run_imaging_pipeline.py;若希望以交互式笔记本方式运行,可在Jupyter中打开run_imaging_pipeline.ipynb。

直接将FITS数据加载到GPU

灵活图像传输系统(FITS)是天文学与天体物理学领域的标准数据格式。通常的做法是先用Astropy在CPU上读取数据,再调用cupy.asarray将得到的数组传输到GPU显存。这意味着解析和解压缩都发生在主机端,之后才通过PCIe进行一次数据传输,而科学计算核函数要等到此时才能真正开始运行。在数据速率达到每小时数TB的场景下,这一整套流程会带来相当可观的数据搬运开销,也是告警延迟预算最先被消耗掉的环节。

为解决这一瓶颈,xDataReader将解压缩、读取和加载FITS数据的过程加速,并直接将数据加载到GPU上。具体而言,xDataReader将整个工作流程拆分如下:CFITSIO在CPU上规划字节范围,KvikIO负责读取(在条件允许时使用NVIDIA GPUDirect Storage),nvCOMP在GPU上对GZIP压缩的图块进行解压,最终由CuPy持有结果数据。该模块支持2到16轴的未压缩图像HDU,以及采用GZIP_1或GZIP_2图块压缩方式的二维图像。目前尚不支持Rice压缩算法以及带抖动处理的浮点量化方式。

在配备兼容驱动与文件系统的环境下,可以使用GPUDirect Storage;在其他CUDA 13机器上,KvikIO可能会走PCIe兼容路径。无论哪种情况,Python接口本身都不会发生变化:batch_to_device依然返回CuPy数组,而cuphoton.xdr.is_gds_active()可以用来查询当前实际使用的是哪条数据路径。

对外公开的API是batch_to_device。它默认读取索引为1的HDU(即第一个扩展区),该区域必须包含受支持的图像数据。如果图像位于主HDU中,则需要传入hdu_indices=(0,)参数。该函数针对每个请求的HDU返回一个堆叠后的CuPy数组。被归入同一次调用中的多个文件,在每个被选中的HDU处必须具有相同的形状和数据类型。

(示例代码部分略,详情参见原始仓库文档)

以下代码会生成两幅合成的256×256帧图像,分别模拟一个参考模板和一次更晚的科学曝光:图中包含三颗静止不动的恒星,以及一个较暗、发生了几个像素位移的运动源。这个唯一的运动源正是经过相减后应当被恢复出来的那个“偶极子”。科学曝光帧的成像也略显模糊,模拟视宁度(大气扰动)较差时的效果。该示例随后对科学图像施加了平移,模拟指向偏差,并更新了头文件信息,使两帧图像仍然描述同一片天区。两个文件并不共享同一套探测器像素网格,但其头文件中的世界坐标系统(WCS)信息,将两者都映射到了相同的天球坐标系中。

对齐与插值:将两次观测对齐到同一套天球网格上

对同一天区的两次观测,通常并不共享同一套像素网格。抖动、旋转和成像比例的差异,会导致同一颗星在模板图像和科学曝光图像中落在不同的像素坐标(x, y)上。FITS头文件中的WCS信息,正是从这些像素坐标映射到天球坐标的桥梁。xPois组件负责在已经完成配准的数组上执行点扩散函数匹配;如果不先完成对齐就直接相减,那么每一颗静止的恒星都会在差异图中产生残差,而运动天体也不再呈现为干净的偶极子图案。

xRep模块可以将二维图像重采样到一套共享的、天球北向朝上的网格坐标系中。build_stack_spec_from_fits函数读取两套WCS解,并构建出一个统一的目标投影区域;reproject_stack函数则将每一帧成员图像都重投影到这套目标网格上。backend="auto"参数会优先选用CuPy,其次是支持CUDA的PyTorch,如果没有可用GPU,则回退到CPU执行。默认的插值方式是Lanczos-3。

接下来这段代码沿用前文生成的FITS文件(即原生探测器像素网格,外加预先设置的CRPIX偏移),将两次观测统一到同一套网格上。完成重投影之后,静止的恒星应当与模板图像中的位置重合,而运动天体则应是唯一一个在天球坐标系中发生了位置变化的源。

(具体代码略)

打印aligned.images.shape和aligned.backend,应当得到一组包含两帧图像的数据堆栈,且在CUDA 13环境下后端应为GPU。如果静止恒星之间仍然存在几个像素的偏差,说明WCS信息并未准确描述你所施加的偏移量,或者目标网格未能完整覆盖两幅图像的投影区域。

PSF匹配与相减

科学曝光图像与模板图像的原始差值,往往会在每一颗亮星周围被残差所主导,这些残差源于大气模糊程度(即视宁度)的差异所导致的PSF不一致。遵循Alard与Lupton提出的方法,最优图像相减并不直接做差,而是求解一个紧凑的卷积核K以及一个差分背景B,使参考模板图像R经过处理后能够获得科学曝光图像的PSF特征,从而与科学曝光T相匹配:

T ≈ R ? K + B

其中?表示卷积运算。随后将这一模型结果从曝光图像T中扣除,得到差异图像D:

D = T – (R ? K + B)

借助基于高斯多项式的核函数基,xPois可以计算出最优图像相减所需要的卷积核。在本例所使用的常数核求解中,backend="auto"参数会依次尝试CuPy、Numba-CUDA、CPU,并选用第一个可用的后端。这一小规模系数求解过程本身留在主机端完成,但像素数据不会离开GPU。采用通量保持的基函数重写方式,可以使核函数的总和保持在接近1的水平;在评估通量保持效果时,应重点检查这一数值。一个表现良好的核函数应当是紧凑且大致对称的;如果出现明亮的环状结构或较大的负值凹陷区域,往往意味着所选基函数无法充分表达视宁度的差异,或者有一颗饱和的恒星未被遮蔽就进入了拟合过程。

接下来这段代码从已对齐的图像堆栈中取出参考图像和目标图像,拟合出一个15×15的核函数。由于xRep已经将两次观测统一到同一套WCS坐标系中,一个紧凑的核函数就应当足以吸收视宁度差异带来的不一致:三颗静止恒星的信号应当被完全抵消,而运动天体则应表现为在t1时刻呈正值的亮斑,在t0时刻呈负值的暗斑。

(具体代码略)

打印fit.backend、kernel.sum()、chi2以及fit_pixel_count等信息。如果核函数总和接近1.0,且残差图中原本围绕亮星的明亮环状结构已经消失,说明匹配过程是成功的。如果残差图中恒星信号仍然占据主导,那么问题可能出在图像尚未完成配准(此时应使用cuphoton.xrep模块将图像统一到同一套WCS坐标系),或者所选核函数基过窄,无法覆盖实际的视宁度差异。在CUDA 13的GPU环境下,fit.backend应显示为cupy,且残差的均方根(RMS)应远低于本场景中约为45的恒星峰值通量。

拟合偶极子

运动天体(例如小行星)以及图像配准误差,都可能在差异图像中产生偶极子图案,即一对配对出现的正、负通量区域。如果对每一个候选目标都单独处理,很容易成为流水线中的瓶颈——每一次曝光,流水线可能都需要在成千上万个小图像切片(stamp)上分别求解最小二乘优化问题。

在本例中,xFit的GaussianDipoleModel将每个图像切片建模为两个旋转椭圆高斯分量之差。这两个分量在同一组像素坐标(x, y)下进行求值,并共享振幅、宽度和方向参数,但各自拥有独立的中心位置(x?, y?)和(x?, y?):

m(x, y) = G(x, y; x?, y?) – G(x, y; x?, y?)

xFit还支持通过StampDipoleModel,基于给定的实采样PSF来构建相应模型。

xFit所采用的Levenberg–Marquardt求解器,可以在CuPy上一次性对大量经过正则化处理的高斯-牛顿方程系统进行批量分解求解。已经收敛的拟合结果会自动停止参与后续的求解迭代,从而减少不必要的计算量。该同步调用会在剩余拟合任务全部完成后,一次性返回整批结果。宽度参数在对数空间中进行优化,并以正的像素标准差形式返回;对于高斯模型,默认采用解析雅可比矩阵。返回的数组(包括参数、协方差和标准误差)均为可移植的NumPy格式,因此分类器环节完全无需导入CuPy求解器。

接下来这段代码会在xPois残差图像中,围绕预先设置的运动天体裁切出一个21×21的图像切片,并从一个略有偏差的初始猜测值出发开始拟合。打印出拟合恢复的中心位置,并与已知的输入位置进行对比,即可检验拟合的准确性。

(具体代码略)

在这组合成数据上,预期应得到converged: True的收敛结果。恢复出的(x_pos, y_pos)和(x_neg, y_neg)应与预先设定的偏移量相差不超过一个像素的若干分之一。如果拟合残差中仍存在结构性的残留信号,通常说明该图像切片中仍包含某颗静止恒星,或者初始中心位置选在了错误的一侧亮斑上。

查看图像切片

上述示例是在一个经过精心控制、仅包含一个已知运动天体、且结果真值已知的合成场景中验证整个工作流的。而真实的观测数据既不具备这样的简洁性,也不具备这样的确定性。单次望远镜指向(即一次观测,visit)就可能产生多达1万个候选目标,一整夜下来更会累积到数百万个,这使得逐一人工检查变得不切实际。因此,下一步工作的重点是对候选目标进行筛选排序,以便优先交由人工复核。

cuphoton xscan review-queue命令可以根据模型不确定性、已知错误类型,或作为纯粹的数据集审查手段,对候选目标进行排序,并限定需要人工复核的候选数量。随后可使用cuphoton xscan review-bokeh命令,查看搜索曝光图像(即当前观测结果)、参考模板图像、原始差异图像,以及xPois残差图像,并为每个候选目标标注“真实”“虚假”或“不确定”。

该复核界面还提供了一个噪声伪影的示例,方便读者在正式用于真实观测数据之前,先熟悉界面布局。由于真实观测得到的FITS数据并未包含在代码仓库中,使用时需要将示例文件复制到examples/xscan/目录下,设置好本地路径,并参照README文档及xScan指南,构建复核队列并启动应用程序。

将同一工作流扩展到多GPU环境

前文演示的内容是在单个GPU上处理一对图像的情形,规模足够小,可以在NVIDIA DGX Spark或工作站上运行。而NVIDIA cuPhoton同样可以在多GPU、多节点系统上运行完全相同的加载与处理路径。完成配准之后,cuPhoton会将完整的图像对及其对应的候选坐标分发到各个GPU工作节点上,并使图像数组在相减、偶极子拟合与分类等整个过程中始终保留在设备端。本示例使用的是合成输入数据和未经训练的模型,仅用于验证多GPU启动流程是否正常;而正是这种可扩展的处理路径,使该工具包能够应对诸如维拉·鲁宾天文台这类巡天级数据流的处理需求。要同时处理多对图像,可参考相关命令;环境搭建、Slurm与SSH启动方式,以及Dragon替代方案的详细说明,均可在分布式执行指南中找到。

(具体命令略)

--images 8参数会生成8对图像,-n 8则会启动8个MPI工作进程,并按一对一的方式映射到各个GPU上。需要注意的是,这些图像对必须已经共享同一套像素网格,也就是前文xRep步骤所完成的工作。如果所使用的机器每个节点只配备一块GPU,则应将两个数字都设置为节点数量,并使用--map-by ppr:1:node参数。上述命令均来自distributed-pipeline示例,其中还包含主机文件模板。需要说明的是,该示例中生成的模型只是用来验证各工作进程是否能够正常启动和结束,并不具备实际的科学意义。

开始使用

以上示例均是在NVIDIA DGX Spark上开发完成的,目标运行环境是配备兼容NVIDIA GPU及驱动的Linux系统。可参照cuPhoton的README文档,克隆代码仓库,检查所支持的Python和CUDA版本,并完成环境搭建。

构建用于FITS输入的原生xDataReader扩展模块:参见xDataReader指南。

在多GPU、多节点环境下运行数据加载与处理:参见分布式执行指南。

创建复核队列并启动候选目标复核界面:参见xScan指南。

了解各组件的API接口与输出产物规范:参见cuPhoton官方文档。

如需了解相关案例,可阅读《纳米尺度成像的加速X射线分析(XANI)》以及《利用加速计算在大型科研设施中实时引导科学实验》等文章。

另外,也可以观看NVIDIA cuPhoton产品概览视频,快速了解这款工具包。

Q&A

Q1:NVIDIA cuPhoton是什么?

A:NVIDIA cuPhoton是一款开源的CUDA-X工具包,提供GPU加速的科学图像处理基础模块,覆盖从传感器数据读取到分类决策的全链路,应用于天文学、X射线科学等高通量数据场景。

Q2:cuPhoton能带来多大的性能提升?

A:在数百TB规模的代表性工作负载上,cuPhoton可将图像加载速度提升最高14900倍,信号处理速度提升最高14550倍,原本需要数月的计算可缩短至数分钟甚至数小时。

Q3:cuPhoton支持哪些硬件和软件环境?

A:cuPhoton 0.1.3版本支持Linux系统下的Python 3.12至3.14以及CUDA 13,可在NVIDIA DGX Spark、单GPU工作站以及多GPU、多节点的Grace Blackwell和Vera Rubin系统上运行。

NVIDIA