TorchMorph：CUDA 加速的数学形态学变换库 TorchMorph: CUDA-accelerated Morphological Transforms
API 逐参数对齐 scipy.ndimage 的 CUDA 形态学、距离变换与最优输运加速库
前置知识
数学形态学（腐蚀与膨胀）
以结构元素在图像上滑动，腐蚀取邻域最小值、膨胀取最大值的一对对偶操作；开运算（先腐蚀后膨胀）去毛刺、闭运算填小洞，形态学梯度与 top-hat 等算子都由它们组合而成。本文中腐蚀/膨胀是同一 CUDA 模板核的两个 functor 实例化，其余 11 个灰度算子全部由这两个原语在主机端派生。
论文的算子分类、核设计（提前退出谓词、边界模式处理）都围绕这对原语展开，不理解它们就看不懂为何两个专属核就能覆盖整个形态学家族。
结构元素与 origin
定义邻域形状的布尔模板（如 $3\times3$ 圆盘），决定每次取极值的范围；origin 指定其锚点相对中心的逐轴偏移。scipy.ndimage 用 size、footprint、structure 三个参数按优先级链指定它。TorchMorph 在主机端把它压平为活跃项列表并预计算扁平偏移。
SciPy 兼容性的核心就是逐项复刻这条参数链和 origin 语义，这是论文 API 设计层与验证矩阵的关键。
距离变换（EDT 与 chamfer）
给每个前景元素标出到最近背景元素的距离。精确欧氏距离变换（EDT）用可分离下包络算法：每轴一遍，每遍计算抛物线族 $f(p)=\min_q\{g(q)+(p-q)^2\}$ 的下包络；chamfer 变换用逐维前后向扫描近似，chessboard/taxicab 分别对应棋盘与曼哈顿度量。
EDT 扫描线的串行性直接决定了核函数的线程块划分方式，也解释了实验中 EDT 批处理收益几乎消失的现象。
熵正则最优传输与 Sinkhorn 迭代
求把分布 $a$ 搬运到 $b$ 的最小代价传输方案，熵正则版在代价上加 KL 项使问题可解：缩放形式迭代 $u\leftarrow a/(Kv)$、$v\leftarrow b/(K^\top u)$，其中 $K=e^{-C/\varepsilon}$；正则很小时缩放形式数值下溢，须转 log 域求解。传输代价对边缘的梯度由对偶势给出。
理解 log 域单 pass 更新、全零边缘映射到 $-\infty$ 的设计以及'对偶势即梯度'的 autograd 实现都需要这一背景。
scipy.ndimage 边界模式
越界邻居的取值约定，灰度核共五种：constant（补 cval）、reflect、nearest、mirror、wrap；二值核只用单一边界值。这些语义是 Python 图像生态的事实标准，下游库被期待完整复现。
论文把'逐参数镜像 SciPy 边界行为'作为核心卖点，差分测试遍历全部边界模式，不熟悉这些模式就无法理解其兼容性验证的意义。
CUDA 融合核与 CUDA Graph
融合核指一次 kernel launch 完成全部计算，而非拼装多个通用张量原语；CUDA Graph 把一串 kernel 捕获成图后整体重放，消除逐次发射的 CPU 延迟。本文 Sinkhorn 每次迭代只含两个小 kernel，长迭代时把 25 次迭代捕获为一块图重放。
这是理解性能来源的关键：$B=1$ 时发射开销占主导，CUDA 图让 1000 次迭代这类长运行从发射瓶颈中解放出来。
研究动机
数学形态学算子在 Python 生态的事实标准 scipy.ndimage 诞生于图像以 NumPy 数组存于主机内存的时代。现代 AI 成像流水线中数据以三种新方式存在：以 PyTorch 张量驻留 GPU、以 batch 一次处理几十个体数据、常见于三维或四维（体积时间序列、多通道断层扫描）。此时每次调用 scipy.ndimage 都要触发设备到主机拷贝、单线程 CPU 计算再拷回，逐样本 Python 循环把天然并行的负载串行化。现有 GPU 库又各有覆盖缺口：Kornia 只支持二维、缺完整边界模式与精确欧氏距离变换；cuCIM 经 CuPy 互操作层而非原生 torch.Tensor；MONAI 把部分形态学后处理交回 SciPy 或 CuPy；OpenCV 与 scikit-image 的快速路径在主机侧且限于二维；熵正则最优传输则住在独立的 POT 栈里。实践者要同时拿到批处理 GPU 形态学、精确距离变换和可微传输距离，必须拼装三个张量约定互不兼容的库。
本文的目标是本文的目标是把上述缺口在一个库里补齐：以原生、批并行的 CUDA 核函数形式，把二值与灰度形态学、精确与近似距离变换、熵正则最优传输共 22 个公开算子带进 PyTorch 张量世界。每个算子都接受形状为 $(B, C, \text{Spatial...})$、空间秩最高为 8 的 CUDA 张量，并在名字、参数顺序、默认值和边界行为上逐参数镜像 scipy.ndimage——包括 structure > footprint > size 优先链、五种边界模式、结构元素 origin、预分配输出缓冲，以及 iterations < 1 表示迭代到结果不变的语义（二值传播和填洞正是靠它构建）——使既有流水线的迁移只是改一行 import 而非重写，例如把逐体积的 ndi.grey_opening 循环换成一次 tm.grey_opening(vols_cuda, size=3) 调用。同时要求三维、四维数据是一等公民而非特例，算子以融合核一次 launch 完成，所有结果与 SciPy、POT 逐元素对齐验证，并以 MIT 许可开源。
与已有工作不同的是，作者的切入角度非常克制：明确声明贡献是'可用性而非算法新颖性'。可分离下包络 EDT、chamfer 扫描、Sinkhorn 迭代这些经典算法都已存在，缺的是有人把它们一次性、正确地按现代成像流水线真正使用的张量约定实现，并用社区已信任的参照系逐元素验证。具体体现在三处别人没做到的地方：其一，SciPy 兼容语义被严格限制在 Python 层，CUDA 层只见归一化几何，这让兼容性定义集中、派生算子零额外设备代码；其二，每次调用只做一次 kernel launch，结构元素几何全部在主机端预先展开为扁平偏移，设备端内部坐标走无边界检查的快路径；其三，验证不是口头断言而是差分测试——78 个测试函数在配备物理 CUDA 设备的自托管 CI 上执行真实核函数而非 mock 路径，与 SciPy/POT 逐元素比对。这种'把老算法做对做快、按新约定交付'的系统工程定位，与追求新模型的论文形成鲜明对比。
核心方法
TorchMorph 的整体思路是'一次调用一次 launch'的融合核设计加三层架构。图 1 把 22 个导出算子组织为结构元素工具组和四大家族：二值形态学（erosion、dilation、opening、closing、propagation、fill_holes、hit_or_miss）、灰度形态学（erosion、dilation、opening、closing、gradient、laplace、white/black tophat）、距离变换（精确欧氏 EDT、chessboard/taxicab 度量 chamfer、作正确性预言的暴力法）与最优传输（SinkhornSolver）。只有加粗条目有专属 CUDA 核，其余是主机端组合，新增 top-hat 变体零设备代码。图 2 展示三层实现：Python 层负责参数校验与结构元素解析，是唯一定义 SciPy 兼容语义的地方；pybind11 绑定层只暴露 8 个核入口、做 dtype 与设备守卫；CUDA 层由 6 个翻译单元组成。库仅依赖 PyTorch 与 CUDA 工具链。
核心创新是围绕'几何在主机端解决、设备端只做纯计算'的融合核设计。朴素 GPU 形态学核要为每个输出元素和每个结构元素位置做完整的 N 维坐标映射与逐轴边界测试：$3^3$ 结构元素配 3-D 输入就是每个体素 27 次映射加逐轴越界检查。TorchMorph 在主机端把结构元素一次性压平为活跃项列表，为每项预存逐轴偏移和相对输入空间步长的扁平偏移，非活跃足迹位置根本不进设备。设备端每个线程先用主机算好的逐轴偏移极值判断自己是否处于内部：占绝大多数的内部线程走快路径，直接把预计算扁平偏移加到线性索引上，无逐轴算术、无边界测试；只有距边界面一个元素内的线程才走一般路径，按 SciPy 语义解析越界邻居——灰度核实现 constant、reflect、nearest、mirror、wrap 五种边界模式，二值核用单一边界值。腐蚀与膨胀是同一模板核的两个 functor 实例化（分别提供单位元与合并规则），二值核再加完成谓词：腐蚀一旦得到假即停、膨胀一旦得到真即停，在稀疏掩码上节省显著。
方法步骤详情
流程分四步。第一步 Python 层解析：按 structure > footprint > size 优先链确定结构元素，把 origin 展开为逐轴元组并做范围检查，把边界模式字符串映射为整数编码，并组合派生算子。第二步 pybind11 绑定层做 dtype 与设备守卫，进入 8 个核入口之一。第三步 CUDA 核执行：形态学按融合方案运行；精确 EDT 用可分离下包络算法，每轴一遍、每遍计算抛物线族 $f(p)=\min_q\{g(q)+(p-q)^2\}$ 的下包络，因扫描串行而每条扫描线一个线程块——单线程在共享内存建包络、整块协同加载、查询阶段并行二搜交点，另有 2-D 特化与全局缓冲回退路径。chamfer 做维度可分离前后向扫描、chessboard 加对角 pass。Sinkhorn 解缩放迭代 $u\leftarrow a/(Kv)$、$v\leftarrow b/(K^\top u)$：每个（行，批 tile）对一个块、tile 含 8 个直方图对；log 域以运行最大值加重缩放和的单 pass 替代两遍 log-sum-exp，空状态良定义的合并算子使全零边缘得 $-\infty$ 而非 NaN；100 次以上迭代按 25 次一块捕获 CUDA 图重放。第四步反向：autograd 返回居中对偶势（由包络定理即传输代价对边缘的精确梯度），反向仅一次广播乘法。
技术新颖性
新颖性不在算法而在系统集成与正确性工程。与 Kornia、cuCIM、MONAI 相比本质区别有三：一是语义保真——border 模式矩阵、结构元素 origin、迭代约定逐项复刻 SciPy，iterations < 1 迭代到收敛正是二值传播和填洞的构建方式；二是核函数层一次性融合——每调用一次 launch、内部快路径消除边界判断，加上 CUDA 图重放（每次迭代两个小 kernel，从 100 次起每 25 次捕获一块）消除长迭代发射延迟，以及批 tiling 使代价矩阵流量按 8 倍摊薄，这些都是把老算法变成训练循环可用组件必需的工程；三是验证文化——78 个测试函数覆盖高秩输入、batch/通道组合、非连续布局、各向异性采样与全部边界模式，参照系逐样本应用于与 GPU 相同的数组以防 batch 掩盖差异，另四分之一是契约测试与运行时测试（侧流、非默认设备）。NaN 传播是唯一记录的与 SciPy 的行为差异。
实验结果
实验在 RTX 4090 D（48GB）上以单线程 SciPy 1.18.0 和 POT 0.9.6.post1 为参照。数值对齐（表 2）：二值形态学与 chamfer 对 scipy.ndimage 逐元素零误差；灰度形态学最大绝对误差 $4.77\times10^{-7}$；精确 EDT $2.06\times10^{-7}$；Sinkhorn 对 POT $1.75\times10^{-6}$，相对 $\ell_2$ 误差均低于 $9\times10^{-8}$。吞吐（表 3）：批执行对单线程 CPU 在灰度形态学上最高 $1.1\times10^3$ 倍、精确 EDT 最高 350 倍；批收益随输入增大而递减——灰度膨胀 $256^2$ 从 $B{=}1$ 的 10.3 到 $B{=}8$ 的 111.1 inputs/ms（10.8 倍）；EDT $1024^2$ 仅 2.11→2.86（1.4 倍）。最优传输（表 4）：scaling 100 次迭代 23.0→1.2ms（18.8 倍），1000 次 229.7→10.1ms（22.8 倍，CUDA 图重放），批量 $n{=}16$ 时 367.3→8.7ms（42.4 倍）。作者注明这是 GPU 对单 CPU 核的对比，暴力法只应读作正确性预言的成本。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 灰度形态学（批量执行） | 加速比（SciPy 时间 ÷ 批量 GPU 单项时间） | 最高 1.1×10³ 倍；grey dilation 256² 达 111.1 inputs/ms（B=8） | scipy.ndimage 单线程（0.756 inputs/ms） | 约 3 个数量级 |
| 精确欧氏距离变换 | 加速比 | 最高 350 倍 | scipy.ndimage distance_transform_edt | 约 2.5 个数量级 |
| Sinkhorn 最优传输（批量 n=16，100 次迭代） | 加速比与耗时 | 42.4 倍（8.7ms vs 367.3ms） | POT 0.9.6.post1 | 约 42 倍 |
| 数值一致性（浮点算子 worst-case） | 最大绝对误差 max|y−ŷ| | 形态学/EDT ≤ 4.77×10⁻⁷；Sinkhorn 对 POT 1.75×10⁻⁶；二值与 chamfer 误差为 0 | scipy.ndimage / POT 逐元素参照 | 达到 float32 精度极限的一致性 |
| 批处理扩展性（grey dilation 256²） | 吞吐 inputs/ms（B=1→B=8） | 10.3→111.1（10.8 倍） | B=1 单输入执行 | 10.8 倍；但 EDT 128³ 恒约 1.4，显示收益随输入增大递减 |
局限与改进
作者承认三项局限：其一，形态学与距离变换核都是纯前向的，只有传输模块可微——不过腐蚀与膨胀可以经由 arg-min/arg-max 位置路由次梯度；其二，形态学与距离核要求 CUDA，预定的回退是 SciPy 本身（传输模块则可回退到纯 torch CPU），CUDA 不可用的环境拿不到批处理收益；其三，核在 --use fast math 下以 float32 计算，NaN 传播不保证与参照一致，这是唯一记录的行为差异。我补充观察：所有加速比都是 GPU 对单 CPU 核（Xeon Gold 6330 单线程）的对比，多线程 CPU 基线会部分缩小差距，1.1×10³ 倍这类数字需谨慎解读；吞吐数据仅来自单张 RTX 4090 D，EDT 这类扫描线串行算子在单输入上已近设备饱和（$128^3$ 批收益恒为零）说明其瓶颈在块间并行度而非发射开销；空间秩硬上限为 8；没有多 GPU、多流性能调度的报告，也未给出 bf16/fp16 等低精度路径的误差特性。
独立分析的弱点
独立分析四个弱点。第一，可微性缺口：边界感知损失（如 boundary loss、Hausdorff 代理）常需对距离场求导，但 EDT/chamfer 前向不可导，用户在训练关键路径上只能退回 torch 原生实现——方向是按作者提示为腐蚀/膨胀实现 argmin/argmax 路由次梯度，并为可分离 EDT 推导基于下包络解析结构的反向。第二，单精度与 fast math：医学影像中近相切掩码场景下 $4.77\times10^{-7}$ 量级误差通常无害，但 NaN 行为不确定会妨碍用 NaN 作哨兵值的代码，建议提供关闭 fast math 的精确模式开关与 float64 路径（目前 float64 仅用于跨实现一致性测试）。第三，算子覆盖面：没有连通分量标记、形态学重建、watershed——这些在分割后处理中与形态学同等常用，作者也列为未来工作。第四，基准公平性：未报告多线程 OpenMP SciPy 对照，1.1×10³ 倍易被误读，建议在文档中显式并列多线程 CPU 基线与 B=1 列，并把暴力法明确标注为正确性预言而非性能项（文中已有此说明但仅一句带过）。
未来方向
作者提出两个方向：为形态学算子加 autograd（经 arg-min/arg-max 位置的次梯度），以及扩充连通分量与重建类算子。基于本文成果可自然延伸：把可分离 EDT 的解析梯度做成可微距离变换，直接服务边界损失与 Hausdorff 代理损失在训练循环中的落地；引入 bf16/fp16 路径并系统量化误差，适配推理侧大批量低成本场景；多 GPU 与多流调度——当前测试已验证侧流与非默认设备的正确性，缺的是性能层面的流并行与设备间分块；把 Sinkhorn 的 CUDA 图重放推广到自适应停止准则（可变迭代数下按最大块捕获）；以及向上集成——以 MONAI/Kornia 插件形式分发，让医疗影像与可微视觉社区无需直接接触 CUDA 工具链。更长远的方向是把同一套'SciPy 语义层 + 归一化几何核层'架构复制到更广的 ndimage 算子面，如 map_coordinates、geometric_transform 等插值与几何变换算子。
复现评估
复现条件非常友好。代码以 MIT 许可开源（https://intcomp.github.io/tm），benchmark/ 目录有逐表复现脚本。基准环境为单机：RTX 4090 D（48GB，driver 580.173.02）、2×Xeon Gold 6330、128GB 内存、Ubuntu 24.04.3、CUDA 12.4、PyTorch 2.6.0+cu124、SciPy 1.18.0、POT 0.9.6.post1，扩展以 -O3 与 --use fast math 编译。计时用 torch.utils.benchmark 阻塞自动量程取中位数、最短 1 秒、计时区内显式 synchronize。正确性由 78 个测试函数保障，每张表都由真实核函数（非 mock）在自托管配物理 CUDA 设备的 CI 上验证。复现门槛是一块 NVIDIA GPU 和 CUDA 工具链；无 GPU 时功能可用 SciPy 回退但性能结论无法复现。论文未提供各内核优化（快路径、CUDA 图、批 tiling）的单独消融，复现单项优化收益需自行实验。
论文图表
两行代码对比：迁移前逐体积调用 CPU 的 ndi.grey_opening(v, size=3) 的 Python 循环，迁移后一次调用 tm.grey_opening(vols_cuda, size=3) 在 GPU 上整批完成，仅一次 launch。
以最直观的方式展示'迁移只需改 import'这一 API 设计目标的达成程度。