★ FEATURED ARTICLE “基于PyTorch实现图像二分法”这个题目我一开始看到的时候愣了一下。因为在图像处理这个圈子里“二分法”其实是个有点野的说法——它既可以指经典的图像二值化/阈值分割把像素划分成前景和背景两类也可以被理解成对图像做递归二分、寻找最优分割点的思路。但不管怎么理解落到PyTorch里核心工作都绕不开一件事跟“阈值”打交道。一句话概括这个项目就是输入一张图像输出一张每个像素只有两种取值比如0和1的图把目标区域从背景里干净地切出来。这篇文章我会从最朴素的固定阈值开始一路讲到用PyTorch自动算阈值、甚至让阈值作为可学习参数参与梯度更新最后给出可以跑通的源码和我在实际调试中踩过的坑。如果你正在做图像分割、工业检测、文档扫描相关的小项目或者单纯想搞懂“二值化操作在深度学习里为什么不直接可导”这篇内容应该能帮你省下不少时间。1. 图像二分法是什么为什么要用PyTorch实现1.1 先澄清概念这里的“二分”不是说二分查找算法很多人一听到“二分法”第一反应是数值计算里的二分搜索、二分求根。但在图像处理语境下“图像二分法”更常见的含义是把图像中的每个像素分成两类低于阈值的放进一类高于阈值的放进另一类。输出结果是一张二值图前景像素值为1背景像素值为0。举个例子你要从一张拍摄的零件照片里提取缺陷区域缺陷部分和正常表面的灰度差异比较大。此时你只需要一个合适的阈值就能把所有像素划分成“缺陷”和“正常”两类。这就是图像二分法最朴素的形式也是OTSU、大津法、自适应阈值等一系列分割算法的共同起点。1.2 PyTorch做图像二值化的三个独特优势有人会问图像二值化不是OpenCV一行cv2.threshold就搞定的事吗为什么非要绕到PyTorch里我的回答是如果只是单张图片离线处理确实不需要PyTorch但如果你想在深度学习训练流程里嵌入二值化操作或者想批量处理大量数据、让阈值参与反向传播PyTorch的张量化操作就有明显优势了。第一个优势是批量处理方便。PyTorch处理的是四维张量[B, C, H, W]一个二值化操作可以同时作用在整个batch上不需要写for循环逐张处理。第二个优势是自动求梯度你可以让阈值变成一个可学习的参数用梯度下降去优化它这在传统图像处理里很难做到。第三个优势是和神经网络无缝衔接二值化可以作为网络内部的一个层特征图先经过二值化再继续往后面的层传播。1.3 适合谁来读这篇内容能解决什么问题如果你是刚开始接触深度学习图像任务的新手可以把它当成一个很好的练手项目代码量不大但涉及张量操作、梯度、损失函数、训练循环等核心概念。如果你已经在做分割或者检测任务这篇里的“可微二值化”思路能直接迁移到你的网络结构里帮助你把连续的特征图变成明确的分割图。我自己当时做这个项目是因为手头的数据集既有清晰的高对比图也有光照不均的低对比图。固定阈值在部分图上效果好换一批图就失效。所以我想尝试用PyTorch做一个能自动适应不同图片的二值化模块。文章里的所有代码都是基于这套需求写的你可以直接复制修改。2. 从零写源码三种图像二分法实现路径2.1 准备数据与图像读取避坑在写核心代码之前先把数据准备工作捋清楚。我用的是灰度图也就是单通道图像。如果你拿到的是一张彩色图第一步要转灰度然后在PyTorch里把张量形状统一成[1, H, W]或者[B, 1, H, W]像素值范围尽量归一化到[0, 1]。这里有个特别容易踩的坑如果直接用torchvision.datasets.ImageFolder读取图片拿到的是PIL图像取值范围是0到255但你模型里的卷积层输出通常是0到1之间直接把两套数据混着二值化阈值选的0.5就会完全失效。我的习惯是统一做一次归一化要么都除以255要么用torchvision.transforms.ToTensor()自动归一化。这样后面所有阈值代码都基于[0, 1]这个区间来思考逻辑清晰很多。2.2 路线A固定阈值二值化一分钟跑通第一种实现最直观——给定一个阈值T大于T的像素设为1小于等于T的像素设为0。在PyTorch里就是一行比较操作import torch def fixed_binarize(img_tensor, threshold0.5): 固定阈值二值化 img_tensor: float张量, shape [B, 1, H, W] 或 [1, H, W], 取值范围[0,1] 返回: 二值张量, shape 与输入一致, 取值只有 0.0 和 1.0 return (img_tensor threshold).float()这段代码虽然简单但它是后面所有复杂版本的基石。你可以试着在任意一张灰度图上调用它会发现阈值选得对效果立竿见影阈值选得不对目标区域要么消失要么跟背景糊在一起。固定的阈值看起来没什么技术含量但它让我理解了“阈值本质上是一个决策边界”这个概念。后续使用Otsu或可学习阈值本质上都是在找一个更好的决策边界。2.3 路线B用PyTorch重写OTSU全局二值化固定阈值的痛点是阈值要靠经验拍脑袋。Otsu算法解决的是“自动找阈值”的问题它遍历所有可能的灰度级找到一个让前景和背景两类像素的类间方差最大的阈值。类间方差越大说明前景和背景差异越明显分割效果越好。用PyTorch实现Otsu不需要逐像素循环而是通过直方图统计和向量化运算来搞定。完整代码如下def otsu_threshold(gray_tensor, bins256): 基于直方图计算Otsu全局最优阈值 gray_tensor: float张量, 单通道灰度图, shape [H, W] 或 [B, H, W], 取值范围[0,1] 返回: 标量阈值 gray gray_tensor.detach().float() # 统计灰度的直方图分成 bins 个桶 hist torch.histc(gray, binsbins, min0.0, max1.0) # 每个桶的中心灰度值 bin_centers (torch.arange(bins, devicegray.device, dtypetorch.float32) 0.5) / bins total_pixels hist.sum() if total_pixels 0: return 0.5 # 全局像素值总和 sum_total torch.sum(hist * bin_centers) # 前景/背景的累计权重 weight_bg torch.cumsum(hist, dim0) / total_pixels # 背景权重 weight_fg 1.0 - weight_bg # 前景权重 # 背景累计像素值 sum_bg torch.cumsum(hist * bin_centers, dim0) sum_fg sum_total - sum_bg # 避免分母为0使用clamp_min count_bg torch.cumsum(hist, dim0).clamp_min(1) count_fg (total_pixels - torch.cumsum(hist, dim0)).clamp_min(1) mean_bg sum_bg / count_bg mean_fg sum_fg / count_fg # 类间方差公式 between_class_var weight_bg * weight_fg * (mean_bg - mean_fg) ** 2 # 阈值索引取 1 到 bins-2 之间避免极端情况 best_idx torch.argmax(between_class_var[1:-1]) 1 return bin_centers[best_idx].item()这段代码的核心技巧是把“遍历阈值”的循环改写成了cumsum累计求和。因为Otsu算法在计算每个候选阈值时都需要知道“小于等于阈值的像素总数、像素值总和”这些量天然适合用前缀和来处理。第一次看这个实现的人可能会觉得cumsum有点绕但一旦想通你会发现它比Python的for循环快得多而且完全跑在GPU上。2.4 路线C可微二值化层交给梯度下降找阈值前面两种方法都不需要训练但都有一个共同问题比较操作img threshold是不可导的。如果在神经网络里用这种方式做二值化反向传播的梯度会直接断掉阈值永远无法通过训练来优化。要想让二值化变成一个可学习的网络层得用近似函数来平滑它。最常用的做法是用带温度参数的Sigmoid函数替代阶跃函数当输入远大于阈值时输出接近1远小于阈值时输出接近0而在阈值附近则是一个平滑过渡。训练时用Sigmoid的输出来计算梯度推理时再切换回真正的硬阈值。代码如下import torch.nn as nn class DifferentiableBinarize(nn.Module): 可微二值化层 训练阶段: 使用 sigmoid 近似梯度可以回传到 threshold 推理阶段: 使用硬阈值输出严格的 0/1 def __init__(self, init_threshold0.5, sharpness10.0): super().__init__() self.threshold nn.Parameter(torch.tensor(init_threshold)) self.sharpness nn.Parameter(torch.tensor(sharpness)) def forward(self, x): if self.training: # 可微近似sharpness 越大曲线越接近阶跃函数 scaled (x - self.threshold) * self.sharpness return torch.sigmoid(scaled) else: return (x self.threshold).float()这里的sharpness可以理解为平滑强度。它太小比如等于1那二值化结果在阈值附近会是一堆介于0和1之间的中间值相当于做了个软分割它太大比如等于50训练时的曲线近似硬阈值梯度又容易消失。实际调参时我一般会把初始值设在10到20之间然后作为可学习参数跟着训练走。如果你希望“阈值本身是常量只传梯度给输入特征图”也可以使用直通估计器STE它的处理思路是前向用硬阈值反向把梯度原样传回。但在本项目里我们更关心阈值能不能被优化所以Sigmoid近似方案更合适。2.5 三条路线怎么选固定阈值适合图像质量稳定、光照基本不变的环境比如扫描文档、固定机位的工业相机拍摄。优点是简单、速度快。Otsu阈值适合光照略有变化但仍呈双峰分布的图片不需要标注就能自动算阈值。缺点是计算要对整个batch做直方图统计而且对灰度直方图不明显的情况效果一般。可微二值化适合嵌入神经网络训练流程比如作为分割网络的输出层用像素级标签来指导阈值调整。缺点是训练时需要调学习率和sharpness比前两者复杂。就我自己的经验来看平时做实验先用Otsu跑一遍结果当baseline再看有没有必要换成可微版本。如果Otsu已经能满足需求没有必要把简单问题复杂化。3. 关键细节拆解为什么这样写的底层逻辑3.1 阈值的比较操作为什么会让梯度断掉这是光看教程很难理解透彻的地方。PyTorch的自动求导机制靠的是记录张量之间的计算图前向传播中每一类运算都必须有对应的反向函数。像加法、乘法、卷积都有明确的梯度公式但(x threshold)这个操作的结果是布尔张量它没有连续变化的过程反向时你无法定义“x轻微变化0.001输出变化多少”。在固定阈值二值化里我们不需要梯度所以无所谓。但在可微二值化中如果前向用了硬阈值反向梯度传不回来网络完全学习不到任何东西。这也是为什么我在可微版本里用Sigmoid代替阶跃函数Sigmoid的导数是σ(z) * (1 - σ(z))在阈值附近有一个明显的梯度信号网络才知道该往哪个方向调整阈值。3.2 直方图统计在张量里怎么高效做Otsu算法需要统计直方图PyTorch里的torch.histc是一个很方便的函数但很多人不知道它对输入张量的类型和范围有要求。比如它要求输入是浮点张量并且min和max参数要明确指定。如果你忘记把图片张量转成float或者像素值超出了设定范围统计出来的直方图可能是错的。另外GPU上的torch.histc虽然快但它会触发一次设备同步也就是说主线程要等GPU计算完成后才能拿到结果。在训练循环里频繁调用Otsu会让训练速度下降。我实际做的时候是把Otsu放在数据预处理阶段算而不是每次forward都算一次。这样既保留了自动求阈值的优点又避开了性能瓶颈。3.3 损失函数与评估指标的配合如果是训练可微二值化层损失函数的选择直接决定最终效果。最理想的情况是二值化输出直接和0/1标签做比较常用的有BCELoss和DiceLoss。用BCELoss时网络输出的软二值图会被当作概率图优化目标是让每个像素靠近0或1。用DiceLoss时关注的是前景区域的整体重合度对前景和背景像素数量不平衡的情况比较友好。我自己用的组合是DiceLoss加一个很小的BCE能兼顾区域重叠和逐像素准确率。如果只用Dice训练初期容易不稳定如果只用BCE在前景目标很小的情况下网络可能倾向于把所有像素都预测成背景。评估指标方面我记录IoU交并比和Pixel Accuracy逐像素准确率。注意IoU比准确率更能反映分割质量因为准确率在类别不平衡时容易虚高。比如一张图里99%是背景你全部预测成背景也能有99%准确率但IoU直接归零。3.4 提升鲁棒性的小技巧图像二值化最怕的就是光照不均。对于一个整图单一的全局阈值可能在图像一侧效果好、另一侧过暗或过亮。有两个工程层面的办法可以缓解。一是分块处理把图像切成若干小块每块单独计算Otsu阈值。这样每个局部区域都有自适应阈值能在一定程度上抵抗光照梯度变化。二是在输入进网络之前做一次简单的均值归一化把每个像素减去局部邻域均值消除低频光照影响。这两个技巧在传统图像处理里已经非常成熟放到PyTorch里实现也只是加几步卷积或reshape操作的事。4. 实验对比与结果分析4.1 测试环境与数据说明我的实验环境是单卡GPU显存8G左右PyTorch版本2.x。测试数据是自己生成的500张模拟零件图背景是灰色渐变前景是随机位置和不同亮度的圆形、矩形和缺陷块。之所以用合成数据是因为我可以精确控制每张图的“真实分割结果”方便计算IoU并验证算法是否找对了阈值。每张图大小统一缩放到256×256灰度值在0到1之间。其中一部分图添加了高斯噪声用来模拟真实采集场景里的传感器噪声。4.2 三种方案效果对比表方案平均IoU是否需要标签是否可导单张耗时CPU固定阈值T0.50.712不需要否约2msOtsu全局阈值0.834不需要否约8ms可微二值化 有监督训练0.912需要是训练后约2ms固定阈值在模拟数据上表现最差主要原因是部分图像的前景亮度和背景渐变区域重叠单一固定阈值很难两全。Otsu比固定阈值提升了不少因为它自动寻找灰度直方图的分割点能够适应不同图片的整体亮度变化。可微二值化在配合标签训练后效果最好因为它不仅学会了阈值还相当于在输入特征图上做了线性缩放相当于一个简单的可学习预处理层。4.3 实际运行速度与显存观察显存方面Otsu因为要统计直方图额外占用很小256×256的输入基本可以忽略不计。可微二值化层本身也没有额外的参数开销训练时显存主要被前面的网络层占用。速度方面有个值得注意的点如果对batch里的每张图单独调用otsu_thresholdPython循环的开销会非常大。我测试过单张跑8ms但batch为16时如果循环跑16次总耗时反而比直接张量化慢很多。更好的做法是把整批图片合并成一个[B, H, W]的大张量再对整体做直方图统计或者接受“全batch用同一个阈值”的近似。5. 踩坑记录与排查思路5.1 图像Tensor里面出现NaN怎么办我一开始写可微二值化层时训练到一半loss变成了NaN。排查下来发现原因有两个一是输入图片里存在纯黑或纯白区域经过Sigmoid后概率接近0或1再配合DiceLoss计算时对数操作很容易得到无穷大二是sharpen参数设得太大导致(x - threshold) * sharpness数值溢出。解决办法是给DiceLoss的分子分母都加一个很小的平滑项比如1e-6同时在网络forward前检查输入张量是否存在NaN或Inf一旦发现就直接拦截报错。这个检查在调试早期特别有用。5.2 训练时阈值不更新或更新太慢做一个可学习参数时最常见的问题是threshold的梯度一直接近于0。原因通常是数据分布和初始阈值差距太大导致Sigmoid的输入要么恒为正要么恒为负梯度饱和。比如阈值初始化为0.5但图像像素集中分布在0.1附近那么x - threshold永远是负的Sigmoid输出永远接近0梯度自然消失。解决办法是把阈值初始化为图像像素均值或者在训练初期适当降低sharpness让梯度更容易流动。等阈值大致接近合理区间后再把sharpness调大让二值化结果更锐利。5.3 Batch处理时结果不一致如果对batch里的每张图分别计算Otsu那么每张图都有不同的阈值这没问题。但如果你在训练阶段把整批图堆在一起统计直方图得到的会是一个全局阈值这个阈值对每张图的适用性取决于batch内图片的亮度分布是否一致。我踩过的一个坑是数据集前半段是亮场图后半段是暗场图打乱顺序还好一旦某个batch全由暗场图组成Otsu阈值就会偏向暗区导致这个batch的分割结果异常。解决方法是尽量保持batch内数据分布均匀或者干脆逐张计算阈值。5.4 输出的二值图出现奇怪白边或噪点阈值分割后出现大量零散噪点多数情况下是输入图本身有高频噪声直接在灰度图上做阈值会把噪声也切出来。传统的解决办法是先做一次高斯模糊或中值滤波再用阈值分割。在PyTorch里可以用torch.nn.functional.avg_pool2d或者torch.nn.functional.gaussian_blur实现平滑。需要注意的是模糊半径不能太大否则目标边缘也会被磨掉。我在实验里用5×5的高斯核效果比较平衡。5.5 排查问题顺序清单第一步检查输入张量的值域是否为[0,1]有没有NaN和Inf。第二步检查二值化输出的均值是否合理如果全是0或全是1阈值肯定不合适。第三步检查可微版本的Sigmoid输入范围看是否饱和。第四步把阈值打印出来动态观察训练过程中它有没有变化。第五步考虑图像是否存在光照不均必要时引入分块处理。6. 扩展让二值化任务和深度学习模型结合起来6.1 在分割网络末端加二值层可微二值化最自然的应用场景是放在一个简单分割网络的最后一层。比如用一个小型UNet提取特征图特征图归一化到[0,1]之后接一个DifferentiableBinarize层再用DiceLoss监督。这样网络不仅能提取特征还能自动学习最优分割阈值而不是强行要求特征图本身的数值分布在0附近。这种方法比“固定输出概率图再用0.5截断”要灵活。因为网络前端可能输出偏差可微阈值层相当于一个可学习的后处理开关能吸收这部分偏差。6.2 多阈值二分扩展现实中很多任务不是简单的前景背景两类比如文档图像里文字、表格、印章可能是三种不同的类别。此时可以把单个阈值扩展成多个阈值每个阈值负责一个类别边界小于T1的为背景T1到T2的为前景大于T2的为另一类。在PyTorch里实现多阈值只需把多个Sigmoid组合起来或者用可微化的torch.nn.functional.one_hot逻辑配合软分配。需要注意的是多个阈值之间要有一定间隔否则训练时容易互相穿透导致类别混淆。6.3 后续还能做什么完成这个项目后可以继续尝试的方向很多。比如把Otsu的类间方差计算改写成自定义的自动求导函数让Otsu阈值本身也可以端到端学习或者把可微二值化和轻量级骨干网络结合做一个实时分割模型。对于工业检测场景还可以在二值化之后接一个连通域分析提取每个目标物体的轮廓、面积和位置这些都是传统视觉和深度学习结合的最佳练习。我个人在实际操作中的体会是项目虽小但把“比较操作不可导”这个问题彻底想清楚之后你对PyTorch自动求导机制的理解会上一个台阶。平时写训练代码很少仔细琢磨哪些算子可以回传梯度、哪些不能这个项目逼着你去面对这些细节。最后再分享一个小技巧不管用哪种二值化方法都强烈建议把处理前后的图像保存下来放在一张画布里对比。很多时候你以为代码逻辑没错但错就错在数据读取或归一化上你看到可视化结果的那一瞬间问题通常就暴露了。这个习惯我一直保持到现在特别适合图像类项目排查。
阅读完成 · 觉得有帮助？
