GPU上跑不动的工具：一个PyTorch扩展如何让形态学处理提速1100倍

你可能没想过这个问题：深度学习模型都跑在GPU上了，为什么处理一张分割掩码、算一个距离场，还得把数据搬回CPU？

这听起来有点荒谬，但确实是很多医学影像和图像分割项目里正在发生的事。你训练一个神经网络分割CT图像里的肿瘤，模型跑在GPU上飞快地吐出预测结果，然后你想对这个分割结果做一次形态学开运算去掉噪点，或者算一下预测边界到真实边界的距离作为损失函数的一部分。这时候你打开Python，习惯性地写下：

```python

import scipy.ndimage as ndi

out = ndi.grey_opening(volume, size=3)

这行代码看起来平平无奇，但它悄悄干了一件很伤性能的事：把GPU显存里的数据整个拷贝到CPU内存，用单线程算完，再拷贝回GPU。如果你的训练循环里每一步都要做这么一次操作，这个数据搬运和单核计算的开销,会像鞋里的一颗小石子,不痛不痒但每一步都硌脚。

这就是这篇论文要解决的问题。作者来自上海大学，做了一个叫TorchMorph的PyTorch扩展库，把形态学处理、距离变换、最优传输这几类经典图像算法,原生地搬到了GPU上,而且做到了和老牌CPU库scipy.ndimage参数对参数地兼容。

问题到底卡在哪里

先说清楚scipy.ndimage是什么。

> scipy.ndimage：Python科学计算生态里处理多维图像的标准库，提供腐蚀、膨胀、距离变换等经典图像形态学算法，是这个领域事实上的行为规范制定者。

这个库设计于图像还老老实实躺在CPU内存里、以NumPy数组形式存在的年代。它没有任何问题，问题是时代变了。

现在的图像处理有三个新特点，而scipy.ndimage一个都没跟上。第一，数据活在GPU显存里，是PyTorch张量，每次调用CPU函数都要经历"显存搬到内存、算完、再搬回显存"这个昂贵的三段式动作。第二，训练是按批次处理的,一次几十个样本一起算,而scipy.ndimage只会处理单张图,你得写个Python循环一张张喂给它,把一个天生可以并行的任务硬生生串行化了。第三，现在很多任务是三维甚至四维的,比如CT体数据、时间序列体数据,CPU算法的耗时会随着空间维度的乘积暴涨。

那有没有现成的GPU图像库能顶上？作者调研了一圈,发现每个候选都有明显的短板。

Kornia支持可微的二维形态学,但不支持高维、不支持scipy.ndimage那一整套边界模式；cuCIM确实是GPU加速的,但它走的是CuPy这条路,不是原生PyTorch张量；MONAI这个医学影像框架,遇到复杂的形态学后处理时,干脆又把活儿甩回给SciPy或CuPy；OpenCV和scikit-image的快速路径基本锁死在二维图像上。至于最优传输,更是活在完全不同的软件栈里,和图像处理库互不来往。

> 最优传输：一种比较两个概率分布"差异有多大"的数学工具,核心思路是计算把一个分布"搬运"成另一个分布所需的最小代价,常用于比较直方图、点云等结构化数据。

所以现实是,一个想要"GPU批量形态学处理+精确距离变换+可微最优传输损失"三件套的工程师,得同时引入三个张量约定各不相同的库,拼凑出一套勉强能用的流水线。这就是TorchMorph要填的坑。

22个算子,一套融合CUDA内核

TorchMorph一共暴露了22个公开算子,分成四大家族：二值形态学、灰度形态学、距离变换、熵正则化最优传输,外加一组结构元素工具函数。

这里有个设计上的取巧之处值得展开讲讲。作者没有给22个算子每个都手写一份CUDA代码,而是只给最核心的几个操作写了专属的融合内核,剩下的算子全部是在这几个原语基础上,用主机端（也就是CPU那一侧的Python逻辑）拼装出来的。

比如二值形态学的开运算、闭运算、传播、填洞,本质上都是腐蚀和膨胀的不同排列组合,论文里说得很直白："只有加粗的条目由专属CUDA内核支撑,其余都是这些原语的主机端组合,这就是为什么新增一种顶帽变种不需要新的设备代码。"

这个决策背后的取舍逻辑,其实和搭积木有点像。如果你造房子的时候,每换一种户型都要重新烧制一批全新形状的砖头,那成本会随着户型数量线性甚至更快地增长。但如果你先烧好几种最基本的标准砖块（长方形、拐角形）,然后靠组合方式去拼出复杂造型,新增一种户型往往只需要重新排列已有的砖块,不需要开新窑。TorchMorph就是把"腐蚀"和"膨胀"当成了那几块标准砖,剩下十几个衍生算子全靠组合搭出来。如果不这样设计,而是每个算子都单独写一份CUDA内核,那22个算子意味着22份需要独立维护、独立测试、独立优化的设备代码,任何一个底层bug都要在22个地方分别排查。

### 融合内核如何省掉重复计算

先说灰度和二值形态学的核心内核设计。

一个朴素的GPU形态学实现会怎么做？对每一个输出像素,都要遍历结构元素覆盖的所有邻居位置,而对每个邻居位置,都要重新把这个像素的一维线性索引拆解成多维坐标,再逐轴做边界检查,判断这个邻居是不是超出了图像范围。

> 结构元素：形态学操作里用来定义"邻域形状"的小模板,比如一个3×3的方块或十字形,决定了腐蚀、膨胀时每个像素要参考周围哪些位置。

论文举了个例子：一个三维图像里用3×3×3的结构元素,意味着每个体素要做27次这样的坐标映射,每次都要沿着每个轴分解索引再做边界判断。这个开销在体数据、批量处理的场景下会被放大很多倍。

TorchMorph用两招把这个开销砍掉了。第一招在主机端提前做：把结构元素展开成一份"有效位置列表",每个位置提前算好它相对于中心的偏移量,以及对应到输入数组存储步长（stride）上的"扁平偏移量"。那些原本就不激活的结构元素位置,压根不会被传到GPU上去,省了无谓的计算。

第二招在设备端做区分对待：每个GPU线程先判断自己负责的这个输出坐标是不是"内部"点,也就是离图像边界足够远,不可能碰到越界问题。如果是内部点（这是绝大多数情况）,直接走一条快速通道：把提前算好的扁平偏移量加到线性索引上就完事,不需要任何逐轴的坐标拆解,也不需要任何边界判断。只有那些真正靠近图像边缘、在结构元素半径范围之内的少数线程,才会走一般路径,去处理各种边界模式（比如常数填充、镜像、环绕等）。

这个"内部快速通道+边界慢速通道"的思路,其实很像高速公路收费站的设计。大部分车流走的是ETC不停车通道,识别车牌、自动扣费,一闪而过；只有少数没装ETC的车,才需要拐进人工收费亭,停车、交钱、找零,慢慢处理。如果所有车,不管有没有装ETC,都被要求统一走人工车道,那高速的通行效率会被那一小撮"边缘情况"拖垮,哪怕它们只占车流的一小部分。TorchMorph正是靠着把"绝大多数是内部点、只有少数是边界点"这个先验知识提前利用起来,才把每个体素27次坐标映射的重活儿,压缩成了大部分情况下一次简单的加法。

腐蚀和膨胀这两个最基础的算子,进一步被实现成了同一份"模板化"内核的两个特例,区别只在于用什么函数去组合邻居的值。二值版本还多加了一个提前退出的判断：一旦某个线程已经能确定最终结果（腐蚀是"只要有一个邻居是背景就判定为背景",膨胀反过来）,就不用再看完剩下的所有邻居了,这在稀疏掩码上能省不少功夫。

### 距离变换：把串行的活儿拆成并行的行

距离变换是这篇论文里另一个重头戏。

> 距离变换：给图像里每一个前景像素标注上它离最近的背景像素有多远,是骨架提取、分水岭分割、边界感知损失函数等一系列应用的基础计算。

TorchMorph提供了精确的欧几里得距离变换、棋盘距离和出租车距离（也叫街区距离）的倒角变换、以及一个用来做正确性校验的暴力算法。

精确欧氏距离变换用的是"可分离下包络"算法,简单说就是把多维的距离计算拆成沿每个坐标轴的一维扫描,每次扫描处理一整条"扫描线"（比如一行像素）,而这条线上的计算本质上是在求一族抛物线的下包络线。

这里有个技术难点：这个下包络的构建过程天生是串行的,你必须按顺序处理这条线上的点,不能像"每个像素独立算"那样随意并行。TorchMorph的解决办法是给每条扫描线分配一整个GPU线程块,块内由一个线程专门负责在共享内存里构建包络,其他线程协作完成数据加载,等包络建好后,再让每个输出位置的查询工作并行地对包络做二分查找。

这个安排就像一个团队做接力报告：写报告大纲这件事必须由一个人从头到尾理清逻辑,没法拆给多人同时写（不然逻辑会乱）,但等大纲定下来之后,团队里每个人都可以同时去查资料、填充各自负责的那一段内容。并行度不是来自"让多个人同时写大纲",而是来自"有足够多份报告需要同时处理"。放到TorchMorph这里,并行度的来源是扫描线的数量：对一个批次B个样本、C个通道、空间维度为n的数据,每一次扫描传播（pass）要处理的扫描线数量是B×C×n的d-1次方,这个数字通常相当可观,足以喂饱GPU的并行能力。

针对不同的数据规模,作者还做了三条不同的执行路径：一条是专门给不超过2048大小的二维图像用的特化版本,把行方向的扫描做成完全连续访存,并且把最后开平方的运算融合进列方向的那一遍扫描里；一条是通用路径,靠转置操作把当前要处理的坐标轴调整到内存里最靠内层的位置,保证每一遍扫描都是连续访存；还有一条兜底路径,当某条扫描线长到共享内存放不下时,把包络栈溢出到全局显存的缓冲区里。

倒角距离变换用的是维度可分离的正向和反向扫描,棋盘距离还额外加了对角线方向的扫描。暴力算法则是把背景点的坐标分批（每批256个）搬进共享内存里,针对不同度量方式和空间维度做模板特化,这样坐标循环能被编译器完全展开,不需要运行时判断走哪种度量。它的时间复杂度是平方级的,所以论文里明确说它的用途不是拿来比性能,而是拿来当"标准答案"验证前面两种快速算法算得对不对——不过有一个例外，倒角变换的接口本身不支持各向异性的采样间隔,所以如果你需要算带各向异性权重的棋盘或出租车距离,只能靠这个暴力版本。

熵正则化最优传输：给Sinkhorn迭代量身定做的批处理

最优传输这一块,TorchMorph走的是熵正则化+Sinkhorn迭代的路线。

> Sinkhorn迭代：求解熵正则化最优传输问题的经典算法,通过反复对一对"势函数"（或者说缩放向量）做行归一化和列归一化,逐步逼近最优的传输方案，得名于提出这个矩阵归一化思路的数学家Sinkhorn。

这个模块的场景设定是这样的：你有n对直方图,维度都是d,而且它们共享同一个d×d的代价矩阵。这个"共享代价矩阵"的特点,给了作者一个明显的优化空间。

一般的批量矩阵向量乘法,会为批次里的每一个样本都重新读一遍那个d×d的代价矩阵,哪怕这个矩阵对所有样本来说都是同一份。TorchMorph的做法是把线程块按照"行,批次分块"的方式组织,每个分块打包八个样本,矩阵的每一行只从显存里读一次,读进来之后同时应用到这八个样本各自的缩放向量上（这些向量就放在寄存器里）。这一下把矩阵的显存访问量除以了八。

这个逻辑很接近一个厨房里备菜的场景。如果你今天要给八桌客人各做一道用同一种酱汁的菜,笨办法是每做一桌就重新去调料架取一次酱汁瓶,来回跑八趟；聪明的做法是一次性把酱汁瓶端到操作台上,八桌菜轮流蘸用同一瓶酱汁,酱汁瓶只搬一次。如果不这样共享,而是让每个样本各自独立地重新加载一遍代价矩阵,显存带宽的浪费会随着批次里样本数量的增加而线性放大,而这原本是完全可以避免的重复劳动。

对数域的更新也做了一处精简。教科书式的log-sum-exp运算通常需要对一行数据扫描两遍：第一遍找最大值,第二遍算减去最大值之后的指数和。TorchMorph维护一个"running maximum"和一个相应重新缩放的"running sum",一遍扫描就搞定,而且这个合并操作在"空状态"（也就是遇到全零边际分布）时也有明确定义,不会输出NaN——直接把对应的势函数钉在负无穷,而不是让计算崩掉。

> NaN：Not a Number的缩写,表示计算中出现的"非法数值"结果,比如零除以零,一旦出现往往会在后续计算里持续传播,污染整个结果。

再往下是一个关于GPU启动开销的观察。作者发现,当迭代次数变多之后,真正的瓶颈往往不是每次迭代本身的计算量,而是每次启动一个GPU kernel所固有的延迟。Sinkhorn的每一步迭代其实只涉及两个很小的kernel调用,如果请求的迭代次数超过100次,TorchMorph会把25次迭代打包成一个CUDA图,之后整体重放这张图,而不是一次次单独发射kernel。

> CUDA图：NVIDIA提供的一种机制,把一系列GPU操作预先"录制"下来打包成一张图,之后可以整体重放,省掉每次单独发射kernel的调度开销。

这就像你如果每天通勤要在同一个红绿灯路口连续经过25次同样的转弯动作,与其每次都重新点击导航软件规划路线,不如提前把这一整套完整路线录成一个"快捷方式",以后直接一键跑完全程,不用每一步都重新问导航"接下来往哪走"。如果不这么做,当迭代次数上到成百上千时,启动开销累积起来会比真正的计算耗时还要拖后腿。而且因为Sinkhorn迭代本质上是在两个乒乓缓冲区之间做不动点更新,多跑几次迭代也不会破坏结果的正确性,所以这种"打包重放"的策略是安全的。

反向传播这块也有个值得一提的设计：梯度的计算靠的是"包络定理"（envelope theorem）,直接返回中心化的对偶势函数,这在数学上被证明就是熵正则化传输代价对边际分布的精确梯度。这意味着反向传播不需要在Sinkhorn的迭代过程里逐步微分,只需要在正向传播时把势函数存下来,反向时做一次简单的广播乘法就够了,省掉了对整个迭代链条求导的巨大开销。

数据说话：到底快了多少

论文的测试机器配置是一块NVIDIA RTX 4090 D显卡,对比的CPU侧用的是单线程的scipy.ndimage和POT库。

> POT：Python Optimal Transport的缩写,是Python生态里另一个专门做最优传输计算的库,这篇论文里用作最优传输部分的性能和正确性对照基准。

先看准确性。二值形态学和倒角距离变换和scipy.ndimage的结果完全一致,连一丝一毫的数值差异都没有。涉及浮点运算的算子,比如灰度形态学、精确欧氏距离变换、Sinkhorn距离,最大绝对误差都不超过1.8×10??,相对误差不超过9×10??。这个级别的误差基本可以归因于GPU用了float32精度以及开启了快速数学优化选项,不是算法逻辑上的偏差。

再看速度提升。灰度膨胀在256×256的图像上,批次从1个样本涨到8个样本时,吞吐量从每毫秒10.31张图涨到111.11张,提升了10.8倍；灰度腐蚀从16.13涨到83.33张每毫秒,提升5.2倍。二值形态学在批量到8时也能做到每毫秒36到56张图不等。

但精确欧氏距离变换的批处理收益就明显没那么夸张了：在1024×1024图像上只从2.11涨到2.86张每毫秒,提升1.4倍；在128立方体数据上几乎是持平的,不管批次大小怎么变都稳定在1.4张每毫秒左右。作者对此的解释是,越是计算量大的重型任务,单个输入自己就已经能占满GPU的并行能力,批处理带来的边际收益自然就变小了。这个规律其实挺直观：一个人搬一块巨石已经用尽全身力气,你让他同时搬两块巨石也快不到哪去；但如果是搬一堆小石子,一次搬一颗和一次搬一把,效率差距就非常明显。

论文摘要里给出的整体数字是：灰度形态学最高提速1100倍,精确欧氏距离变换最高提速350倍,Sinkhorn求解器最高比POT快42倍。这几个"最高"都是相对单线程CPU参照系而言的,论文自己也很坦诚地指出了这一点：这些对比是"一块GPU对一个CPU核心",不是和一个充分并行优化过的多核CPU实现比较,读者应该把这个当成设备能力差异加批处理效应的综合结果,而不是纯算法层面的胜利。

最优传输部分,在32×32的网格上,标准缩放形式跑100次迭代提速18.8倍,跑1000次迭代（用上CUDA图重放）提速22.8倍；对数域形式在小正则化系数下（这时候标准缩放形式容易发生数值下溢）提速11.8倍；而在批次数为16的场景下,提速达到42.4倍,这也印证了前面说的"共享代价矩阵、批量分块"的设计确实带来了实打实的收益。

三层架构：把SciPy的规矩和CUDA的效率分开放

TorchMorph整个代码库分成三层。

最上面是Python层,负责把用户传进来的参数按照scipy.ndimage的规矩解读和校验,比如从structure、footprint、size这三种可能的输入里按优先级挑一个来确定结构元素的具体形状,把origin参数展开成每个维度对应的偏移量并检查是否越界,把字符串形式的边界模式（比如"reflect"、"wrap"）翻译成整数编码传给底层。这一层刻意被设计成唯一"懂"scipy.ndimage参数习惯的地方。

中间是绑定层,用pybind11这个工具把8个CUDA kernel的入口暴露给Python侧,顺带做数据类型和设备的合法性检查。

最下面是CUDA层,由6个编译单元构成,真正的计算逻辑都在这里。

这个三层拆分带来一个直接的好处：CUDA层完全不需要知道scipy.ndimage的参数是什么规矩,它看到的永远是已经被规范化过的、干净的几何信息。这就像一个餐厅的运作分工：前台服务员负责听懂顾客的各种口味要求、忌口、加辣程度这些五花八门的表达方式,把它们翻译成后厨能看懂的标准工单；后厨的厨师只管按照标准工单炒菜,不需要每道菜都重新去猜顾客到底是什么意思。如果这两层职责混在一起,后厨的每一道菜逻辑里都要嵌入"如果顾客说的是这个词就这样做,如果说的是另一个词就那样做"的分支判断,代码会变得又乱又难维护，而且CUDA代码里塞进大量字符串解析和条件分支,本身对GPU的执行效率就是一种拖累。

测试怎么做的

TorchMorph不是靠嘴说"我和scipy.ndimage一样",而是靠78个测试函数,分布在5个测试模块里,做逐元素的差分测试来验证。

测试覆盖了二维、三维乃至更高维度的输入、各种批次和通道组合、非连续和转置过的内存布局、各向异性采样、每一种边界模式、以及带偏移原点的非对称结构元素。对每一种配置,都拿同一份输入数组分别喂给GPU实现和CPU参考实现,逐元素比较,这样即便是批处理场景,也不会出现"整体平均误差看起来小,但某个具体样本算错了"这种被掩盖掉的问题。

对于没有现成参考答案的情况,测试转而验证一些必须成立的数学性质,比如最优传输算出来的传输方案必须能还原出原本的两个边际分布，距离变换返回的"最近背景点索引"必须真的指向一个背景像素，梯度的解析解要用有限差分的方法做交叉验证。

除此之外还有一批"接缝测试"，专门检查不同实现路径之间是否会打架：融合CUDA内核算出来的结果要和纯PyTorch写的兜底实现一致，float32和float64的结果要一致，CPU和GPU的结果要一致，CUDA图重放模式和普通逐次启动模式的结果要一致，提前退出的版本和完整迭代到底的版本要一致。还有一批"契约测试"，专门验证系统在遇到非法输入时是不是老老实实报错，比如空间维度超过8、结构元素原点设置不合理、传了不认识的边界模式字符串，或者传进来的根本不是CUDA张量。

这一整套测试在每次代码提交时都会跑在一台配备了真实物理GPU的自建持续集成服务器上,这意味着每一次改动都是拿真实的kernel去和CPU参考实现对比,而不是靠一个被mock出来的假环境蒙混过关。

局限性,作者自己也说了

论文在结尾很坦率地列了三条限制。

第一,形态学和距离变换的内核目前都是只能前向计算,不支持反向传播,只有最优传输模块是完全可微的。不过腐蚀和膨胀这两个操作,理论上是可以定义次梯度的（把梯度路由到取最小值或最大值的那个位置上）,这部分留给未来工作。

第二,这个库依赖CUDA环境,没有GPU的话,预期的退路是直接改用SciPy本身,而最优传输模块倒是准备了一个纯PyTorch的CPU回退路径。

第三,计算全程用float32精度,并且开启了--use_fast_math这个编译选项来换取速度,这意味着NaN的传播行为不保证和SciPy完全一致。这一点论文里专门强调了是"唯一有文档记录的、和SciPy行为不同的地方"。

未来的方向,作者提到想给形态学算子加上自动微分支持,并且扩展出更多连通域分析和形态学重建相关的算子。

写在后面

读完这篇论文,我最触动的地方其实不是那些提速几百倍的数字,而是作者选择的这个切入角度：不去发明新算法,而是把已经存在了几十年的经典算法,原原本本地按照社区已经信任的行为规范,重新实现一遍。

这在学术圈里是一件容易被低估的工作。发一篇论文,提出一个全新的模型架构,听起来更"有创新性"；而"把scipy.ndimage的距离变换在GPU上原样实现一遍,连边界模式的坑都要一个个踩过去对齐",这种工作往往不容易被当成"研究成果",但它解决的是一个真实存在、每天都有人在吃苦头的工程痛点。

论文里那句"22个公开算子只有8个真正对应CUDA kernel入口"也让我想了挺久。这种"少写核心代码,靠组合覆盖需求"的策略,本质上是在做一种取舍：牺牲了每个衍生算子单独优化的可能性,换来了整个代码库的可维护性和可信度。对于一个想要长期维护、被社区信任的基础设施库来说,这个取舍大概率是值得的。

一个开放的问题是：如果未来真的要给腐蚀膨胀这些操作补上完整的自动微分支持,次梯度这种"退而求其次"的近似方案,在真实的分割训练任务里到底会带来多大的性能损失？这大概是接下来最值得追问的地方。

Q&A

Q1：TorchMorph是什么？

A：TorchMorph是上海大学团队开发的一个开源PyTorch扩展库，用CUDA原生实现了22个图像处理算子，覆盖二值形态学、灰度形态学、精确和近似距离变换、以及熵正则化最优传输，API参数完全对齐scipy.ndimage，可以直接在GPU上批量处理，不需要来回搬数据到CPU。

Q2：TorchMorph比scipy.ndimage快多少？

A：在单线程CPU对比下，灰度形态学最高提速约1100倍，精确欧几里得距离变换最高提速350倍，最优传输的Sinkhorn求解器比POT库最高快42倍，具体提速幅度会随批次大小和数据规模变化，计算量越大的任务批处理收益相对越小。

Q3：TorchMorph的计算结果准确吗？

A：二值形态学和倒角距离变换的结果和scipy.ndimage完全一致，涉及浮点运算的算子（如灰度形态学、精确欧氏距离变换）最大绝对误差不超过1.8×10??，误差主要来自float32精度和快速数学优化选项，唯一已知的行为差异是NaN值的传播方式可能和SciPy不完全一致。