做语义分割的都知道一张图里真正让你头疼的往往不是局部的边缘而是“全局关系”。比如一个路灯杆近看是细长条模型完全能识别可一旦路灯和远处的白色墙面连在一起特征图上几个像素共享同一块感受野网络就开始犯迷糊要么把路灯并进墙里要么把墙跟天空混成一片。这时候让网络“退后一步看全景”比堆更多卷积更有效。PSP模块Pyramid Scene Parsing 模块就是干这件事的用金字塔池化把特征图分成 1×1、2×2、3×3、6×6 的网格分别聚合不同粒度的上下文再上采样拼回去。很多分割网络的涨点都靠它。这篇文章是我在分别用 PyTorch 和 TensorFlow 复现这个模块时做的小结。全文不聊论文废话直接给能跑的代码、形状变化以及两个框架之间最容易踩的坑——尺寸整除、插值对齐、BN 的 training 开关。如果你已经能调通 U-Net想给网络加一个上下文模块或者正在把 PyTorch 模型迁到 TF 部署这篇应该能帮你省不少时间。1. PSP模块的设计思路与核心原理1.1 为什么分割任务需要额外上下文普通卷积堆叠出来的感受野是“慢慢长”的。一个 3×3 卷积经过十几层之后理论上感受野能覆盖全图但实际有效感受野远没有理论值那么大中心区域占了绝大部分权重。这个现象在分割任务里非常致命远处一整片相似纹理可能被归成同一类单个小目标也可能因为周围环境相似而被吞掉。我常用一个生活类比来解释看一张几百人的合影如果只看每个人的脸你很难判断这些人站在哪里但如果你先看整体构图再看局部细节就能同时知道“这是个大场面”和“这一排人正在做什么”。PSP模块要解决的就是让网络同时拥有这两种“焦距”。它不靠增加卷积层数而是直接把特征图切成不同分辨率的网格让网络同时看到全图统计信息、区域统计信息和局部细节信息。这个思路和全局平均池化不一样。全局平均池化是把整张图压成一个点信息保留程度太低只靠多尺度卷积又会在计算量和参数上付出很大代价。金字塔池化正好折中网格越粗覆盖范围越大语义越全局网格越细越接近原始位置空间细节越完整。1.2 金字塔池化的设计逻辑PSP模块的输入通常是一张特征图形状记作 H×W×C。模块内部会并行做四路池化目标网格尺寸分别是 1×1、2×2、3×3、6×6。以 1×1 为例就是全局平均池化把整个特征图压缩成单个向量代表“整幅画面到底属于什么场景”。2×2 把画面分成四块每块告诉网络“左上角是天空、右下角是路面”。3×3 和 6×6 进一步细化保留更多空间对应关系。每一路池化之后特征图的尺寸会变小比如 24×24 的特征图经过 6×6 池化后变成 6×6。这时候如果直接拼接回 24×24尺寸对不上所以需要把每个池化结果做 1×1 卷积降维通常降到输入通道数的四分之一然后双线性插值上采样恢复成 H×W。最后把原始特征图和四路上采样结果按通道维度拼接再接一个 3×3 卷积做融合。维度上有一个容易忽略的点如果输入 C2048四路分支每路降到 512拼接后就是 2048 4×512 4096。这个通道数非常大所以最终融合卷积一般会再降到 512 或 256否则后续分类卷积计算量会失控。很多人在复现时只关注池化部分忽略了通道压缩导致显存直接爆掉。1.3 模块在常见分割网络中的位置PSP模块不是一个独立网络而是一个可插拔的积木。最经典的位置是放在主干网络之后、分类卷积之前。主干网络可以是 ResNet 这类通用分类网络最后几层去掉全局池化和全连接保留带空间分辨率的特征图PSP模块接在后面做多尺度上下文聚合最后再接一个 1×1 卷积输出每像素的类别概率。除了原版的 PSPNet这个模块也能塞进 U-Net 的 bottleneck、DeepLab 的 decoder 之前甚至放在实时分割网络的轻量主干后面。和 DeepLab 系列常用的 ASPP 相比PSP 更依赖池化操作而不是空洞卷积参数更少部署时对框架的算子要求也更低。正因为这种通用性我才特意在两个主流框架里各写了一份实现方便在不同项目里直接抄。2. PyTorch实现要点与代码2.1 选择 AdaptiveAvgPool2d 的原因PyTorch 里实现 PSP 模块最省心的地方就是torch.nn.functional.adaptive_avg_pool2d这个函数。它的含义很直接不管输入特征图是多大我都让你输出一个固定尺寸的网格。输入是 24×24目标 6×6它就自动把 24×24 分成 6 块区间输入换成 31×29它也能算不需要你手动去凑整。这个特性在实际工程里非常值钱。分割训练的输入尺寸经常不固定或者会在训练和推理时使用不同的分辨率。如果用普通的avg_pool2d必须自己算 kernel size 和 stride输入一变就得跟着改代码。用 AdaptiveAvgPool2d 就没有这个问题这也是 PyTorch 社区复现 PSP 模块几乎都这么写的原因。需要注意 AdaptiveAvgPool2d 的窗口划分和普通平均池化不完全一样它对边界块的处理是取整的每块覆盖的像素数可能略有差异。不过这种差异对语义分割这种任务来说基本没有影响因为后续还有 1×1 卷积和上采样特征图信息已经被高度抽象了。2.2 PyTorch 完整实现代码下面是我在项目里常用的 PSP 模块实现保留了足够多的注释方便对照维度变化。import torch import torch.nn as nn import torch.nn.functional as F class PyramidPoolingModule(nn.Module): def __init__(self, in_channels, out_channels512, bin_sizes(1, 2, 3, 6)): super().__init__() self.bin_sizes bin_sizes # 每个池化分支用 1x1 卷积压缩通道降到输入通道的 1/len(bin_sizes) inter_channels in_channels // len(bin_sizes) self.reduce_convs nn.ModuleList([ nn.Sequential( nn.Conv2d(in_channels, inter_channels, kernel_size1, biasFalse), nn.BatchNorm2d(inter_channels), nn.ReLU(inplaceTrue) ) for _ in bin_sizes ]) # 拼接后的融合卷积 self.fusion nn.Sequential( nn.Conv2d(in_channels len(bin_sizes) * inter_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): H, W x.shape[-2:] features [x] for size, conv in zip(self.bin_sizes, self.reduce_convs): pooled F.adaptive_avg_pool2d(x, output_sizesize) pooled conv(pooled) # 上采样回原始分辨率align_cornersFalse 与 TF 的 half_pixel_centersTrue 对齐 up F.interpolate(pooled, size(H, W), modebilinear, align_cornersFalse) features.append(up) concat torch.cat(features, dim1) return self.fusion(concat)这段代码的 forward 流程一目了然先取原始特征图作为第一路然后四路池化、卷积、上采样最后拼接融合。输入输出形状只差通道数空间尺寸保持不变所以后面可以直接接分类头。2.3 通道设计与网络组合通道设计上最常见的一组配置是主干输出 2048 通道四个分支每路降到 512拼接后 4096再通过融合卷积降到 512。这样 PSP 模块输出的通道数就和 DeepLab 的 ASPP 输出接近后续接 256 通道的分类头也比较顺。如果主干输出通道比较少比如只有 320我建议把inter_channels改成in_channels // len(bin_sizes)也就是每路 80拼接后 640融合卷积输出 256。这样参数不会爆炸MobileNet 这类轻量主干也能撑得住。在 PyTorch 里接入主干很简单定义一个包含 PPM 和分类卷积的分割头即可。class PSPHead(nn.Module): def __init__(self, in_channels, num_classes, mid_channels512): super().__init__() self.ppm PyramidPoolingModule(in_channels, mid_channels) self.cls nn.Conv2d(mid_channels, num_classes, kernel_size1) def forward(self, x): x self.ppm(x) return self.cls(x)实际训练时主干和分割头可以分别初始化。主干用分类数据集上预训练好的权重PSPHead 里的参数随机初始化然后一起做端到端微调。我试过先冻结主干只训练 PSPHead等损失降下来再解冻效果比一上来就全参数训练更稳。3. TensorFlow实现要点与代码3.1 TensorFlow 生态下先想清楚三件事写 TensorFlow 版之前有几件事必须提前想清楚否则代码写着写着就会卡住。第一TensorFlow 默认的数据格式是 channels_last也就是张量形状是 B×H×W×C。PyTorch 是 channels_first习惯上 B×C×H×W。如果你平时 PyTorch 写得多刚切到 TF 时会非常容易在维度操作上报错尤其是tf.concat和tf.image.resize的轴参数。第二TensorFlow 没有一个像 PyTorchAdaptiveAvgPool2d那样“一行代码输出任意尺寸”的常用高维接口。它的tf.nn.avg_pool需要你给定 kernel size 和 stride窗口大小得自己算。所以要么限制输入尺寸能被目标网格整除要么像我下面这样做一个简单的工程近似。第三建议直接用 Keras 函数式 API 搭模块而不是把池化和上采样都塞进自定义 Layer。因为tf.image.resize的目标尺寸可以是一个动态张量而 Keras 的UpSampling2D层要求静态size两者在动态输入下行为差很多。3.2 TensorFlow 完整实现代码下面是 TensorFlow 2.x 里的函数式实现。为了降低复现门槛我用一个自定义的adaptive_avg_pool_2d完成“任意目标尺寸”的近似先把输入缩放到目标尺寸的整数倍保证能被整除再做平均池化。这个方案在动态输入下很稳定。import tensorflow as tf from tensorflow.keras import layers def adaptive_avg_pool_2d(x, output_size): 近似 AdaptiveAvgPool2d。 先把特征图缩放到 output_size 的整数倍再做平均池化。 B, H, W, C tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2], tf.shape(x)[3] new_h output_size * tf.maximum(1, H // output_size) new_w output_size * tf.maximum(1, W // output_size) resized tf.image.resize(x, (new_h, new_w), methodbilinear, half_pixel_centersTrue) pool_h new_h // output_size pool_w new_w // output_size pooled tf.nn.avg_pool(resized, ksize[1, pool_h, pool_w, 1], strides[1, pool_h, pool_w, 1], paddingVALID) return pooled搭 PSP 模块时我把每一路的卷积和 BN 也放在函数里用函数式 API 串起来。def psp_module(input_tensor, out_channels512, bin_sizes(1, 2, 3, 6)): x input_tensor H tf.shape(x)[1] W tf.shape(x)[2] in_channels x.shape[-1] inter_channels in_channels // len(bin_sizes) branches [x] for size in bin_sizes: pooled adaptive_avg_pool_2d(x, size) pooled layers.Conv2D(inter_channels, 1, paddingsame, use_biasFalse)(pooled) pooled layers.BatchNormalization()(pooled) pooled layers.ReLU()(pooled) up tf.image.resize(pooled, (H, W), methodbilinear, half_pixel_centersTrue) branches.append(up) concat layers.Concatenate(axis-1)(branches) out layers.Conv2D(out_channels, 3, paddingsame, use_biasFalse)(concat) out layers.BatchNormalization()(out) out layers.ReLU()(out) return out这段代码和 PyTorch 版逻辑完全一致。注意psp_module只是一个搭模型的辅助函数只能在构建模型时调用一次不要在训练循环里反复调用否则每次都会创建一批新的卷积层。3.3 与预训练 backbone 集成接主干时我用 Keras 函数式 API 做一个完整的分割模型。主干可以选择带空洞卷积的版本保证输出分辨率不至于太小。backbone create_backbone() # 输出形状 B x H/16 x W/16 x C基于通用分类数据集预训练 ppm_out psp_module(backbone.output, out_channels256) head layers.Conv2D(num_classes, 1, nameseg_head)(ppm_out) model tf.keras.Model(backbone.input, head)BN 层在 TF 里由 Keras 自动管理训练时用model.fit会更新均值和方差推理时自动切到使用累计统计量。这一点比 PyTorch 手动model.train()/model.eval()稍微省心一点但有一个陷阱如果先冻结主干BN 依然会用累积均值做推理不会因为主干trainableFalse就停止工作。反而是在微调阶段如果主干 BN 被设置为可训练而分支 BN 也保持默认状态容易出现细微差异。我的习惯是微调阶段主干的 BN 保持trainableFalse只训练 PSP 模块里的 BN。4. 实操过程把小模块搭成分割头4.1 输入尺寸与数据准备在我自己跑过的项目里PyTorch 版对输入尺寸几乎零要求任意短边都能跑。但 TensorFlow 版为了配合我上面写的简化池化最好把训练尺寸定在一个能被 2 和 3 同时整除的数值上。以输出 stride 为 16 的 backbone 为例输入 384×384 时特征图是 24×24能被 1、2、3、6 都整除输入 480×480 时是 30×30同样能整除。这两个尺寸实测下来效果都不错。数据增强方面我一般用随机缩放 随机裁剪 左右翻转。裁剪尺寸会直接决定特征图是否能被池化网格整除所以如果你换了训练尺寸建议先打印一次特征图形状确认。不要想当然这个问题在 TF 版里真的很容易踩。4.2 训练配置与损失训练配置上主干一般加载分类预训练权重PSP模块随机初始化。损失函数用普通的交叉熵就能跑出不错的效果如果类别不均衡可以再加一个带难例挖掘的交叉熵。PSPNet 还有一个经典习惯在主干的中间层加一个辅助监督辅助损失权重通常在 0.4 左右。这个辅助头不是必须的但我实测加上之后早期训练更稳最终指标会高一点点。优化器我习惯用带动量的 SGD初始学习率 0.01配合 poly 学习率衰减策略。在很多分割数据集上这个组合比 Adam 更稳。如果是小数据集或者只想验证模块能不能跑通用 Adam 也可以但学习率要调低一些比如 0.001否则 BN 很容易在前期震荡。4.3 推理与后处理推理阶段模型输出是原始输入尺寸的 1/16需要双线性插值上采样回原图再做 softmax 和 argmax。多尺度推理是个便宜又有效的涨点手段常见做法是把输入缩放成 0.75、1.0、1.25 三个尺度分别预测再把 softmax 概率平均。加了 PSP 模块之后多尺度的收益依然存在因为全局上下文对不同尺度的响应会有差异。滑动窗口推理对 PSP 模块来说反而没那么合适。原因是窗口变小之后1×1 分支只在窗口内做全局池化会丢失整个图像的上下文信息。如果输入图特别大我更建议把整图等比缩放到一个合适尺寸做一次全局推理再用双线性插值把预测图放大回原图。4.4 一个小实验的直观观察我做过一个对比实验同一个 backbone同一个训练配置只差有没有 PSP 模块。加了之后最明显的变化是那些“大块物体边界处的零散误判”少了很多。比如一大片草地上有一小片裸露泥土原本只靠局部特征很容易把泥土和草地混在一起加了 PSP 之后网络能利用 1×1 和 2×2 分支的全局统计知道这块区域大概率属于草地中的异常从而更容易保留边界。另一个变化是训练曲线稳定得更快。PSP 模块相当于给网络提供了一个全局统计的“锚点”前几个 epoch 的 loss 下降速度比不加模块时明显更快。这个现象在轻量主干上尤其明显因为轻量主干本身的感受野不够池化分支正好补上了短板。5. 常见问题与排查技巧实录5.1 池化网格与特征图尺寸不整除我遇到最多的坑是 TensorFlow 端池化输出尺寸不对。举例来说特征图是 14×14目标池化到 6×6如果直接按 14//62 作为 stride 做平均池化输出会是 7×7不是 6×6。第一次跑的时候我甚至没留意到 shape 变了直到后面拼接时报维度错误才知道出了问题。解决办法有三个。第一训练尺寸固定成能被所有 bin_sizes 整除的值比如 384×384 搭配 stride 16这样最简单的窗口池化就够用。第二改用我上面写的adaptive_avg_pool_2d近似方案先缩放再做池化。第三动态输入下追求严格一致就要用循环把每个输出格子对应的输入区间算出来再对每个区间分别取均值。这个第三种方案代码量会大不少我一般只在做论文复现时才写。5.2 双线性插值的对齐差异PyTorch 的F.interpolate和 TensorFlow 的tf.image.resize表面看起来都是双线性插值但实际上像素坐标系定义不同。PyTorch 里align_cornersFalse对应坐标中点对齐TensorFlow 里要设置half_pixel_centersTrue才能得到一样的插值结果。反过来PyTorch 的align_cornersTrue对应 TensorFlow 默认的half_pixel_centersFalse。PyTorchTensorFlowalign_cornersFalsehalf_pixel_centersTruealign_cornersTruehalf_pixel_centersFalse这个差异在对齐两个框架的输出时特别重要。我一开始在 TF 端漏了half_pixel_centersTrue跑出来的特征图总是比 PyTorch 版在边缘处差一点点虽然不影响最终分割指标但会让模型移植时非常难排查。5.3 BatchNorm 在 Keras 里的 training 开关在自定义 Layer 里手写 PSP 模块时BatchNorm 的training参数很容易被忽略。如果你用的是 Keras 函数式 APImodel.fit和model.predict会自动切换训练/推理状态问题不大。但如果你把模块封装成自定义 Layer 并在call里手动调用子层就一定要加上training参数并传给 BatchNormalization否则验证阶段会拿着训练批次的统计量做推理结果波动很大。另外还有一个迁移学习的老坑给主干设置trainableFalse时Keras 会冻结全部权重但如果 BatchNormalization 的trainableFalse它用的是 moving_mean 和 moving_variance而不是当前 batch 的统计量。这对迁移学习是合理的但如果你希望微调阶段更新 BN 的统计量就不能简单地把所有层都设为trainableFalse需要单独处理 BN 层。5.4 显存占用与轻量化PSP 模块四路并行池化看起来会占不少显存实际上因为每个分支都先压通道再上采样显存增长并不夸张。真正吃显存的是拼接后的 4096 通道特征图。如果显存吃紧我第一推荐减少inter_channels比如从in_channels // 4降到in_channels // 8拼接通道数会立刻变小。第二推荐去掉 6×6 分支因为它在计算量和参数量上都偏高但对轻量主干来说收益有限。移动端部署时我更倾向于只用 1×1、2×2、3×3 三个分支并且把 fusion 卷积从 3×3 改成 1×1。这样 PSP 模块基本就是四五个 1×1 卷积加几次池化非常轻量对算力敏感的芯片也很友好。6. 选型建议与后续扩展6.1 两个框架输出是否一致如果你只是想验证模块实现得对不对可以在两个框架里分别搭一个随机初始化的 PSP 模块喂同样的输入只比较输出 shape 和数值量级。由于卷积核初始化不同数值不可能完全一致但形状必须相同且每个位置的数值应该在同一个分布范围内。如果要精确对比需要先把 PyTorch 的权重转换到 TensorFlow注意卷积核的排布方式不一样。PyTorch 的 Conv2d 权重形状是 [out, in, kh, kw]TensorFlow 的 Conv2D kernel 是 [kh, kw, in, out]转换时要 transpose。我在实际项目里转换过加上插值对齐之后模块输出的最大绝对误差通常能压到 1e-3 以下。6.2 什么时候选 PyTorch什么时候选 TensorFlow使用场景推荐框架理由论文复现、快速改结构PyTorchAdaptiveAvgPool2d更省心调试直观已有 TensorFlow/Keras 生态TensorFlow与现有训练推理链路无缝衔接移动端 / TFLite 部署TensorFlow量化工具链更成熟转换路径短大量分布式训练两者皆可取决于团队现有基础设施从我自己的经验看如果是新项目从零开始我倾向 PyTorch因为调试体验好写 PSP 这种动态尺寸模块几乎不用考虑对齐问题。但如果项目最后要落到移动端或者 TF Serving那我会直接用 TensorFlow 写省去两套框架之间转换权重的麻烦。6.3 模块的变体与扩展方向PSP 模块的下一步玩法其实很多。最简单的变体是调整 bin_sizes比如去掉 6×6或者改成 1、2、4不同数据集上最优组合不一样。还有一种思路是把池化后的特征经过 1×1 卷积后加一个注意力权重让网络自己决定哪个尺度的上下文更重要。这相当于把简单的拼接变成了加权融合涨点空间更大。另外PSP 和 FPN 这类特征金字塔结构并不冲突。可以把 FPN 各层的输出分别过 PSP 模块再逐层融合这样每个尺度的特征都有全局上下文兜底。代价是计算量上去了但效果通常也比单层加 PSP 更稳。我一般先不加只有当单层模块提升不明显时才往多尺度方向扩展。这几年每次接到分割任务我第一反应不是去换整个网络结构而是先往 baseline 里塞一个金字塔池化模块。它不一定是涨点最猛的却是我用过最稳的上下文模块之一。如果你也在 PyTorch 和 TensorFlow 之间来回折腾希望这篇小结能帮你把两个框架里的坑一次性趟平。
做语义分割的都知道一张图里真正让你头疼的往往不是局部的边缘而是“全局关系”。比如一个路灯杆近看是细长条模型完全能识别可一旦路灯和远处的白色墙面连在一起特征图上几个像素共享同一块感受野网络就开始犯迷糊要么把路灯并进墙里要么把墙跟天空混成一片。这时候让网络“退后一步看全景”比堆更多卷积更有效。PSP模块Pyramid Scene Parsing 模块就是干这件事的用金字塔池化把特征图分成 1×1、2×2、3×3、6×6 的网格分别聚合不同粒度的上下文再上采样拼回去。很多分割网络的涨点都靠它。这篇文章是我在分别用 PyTorch 和 TensorFlow 复现这个模块时做的小结。全文不聊论文废话直接给能跑的代码、形状变化以及两个框架之间最容易踩的坑——尺寸整除、插值对齐、BN 的 training 开关。如果你已经能调通 U-Net想给网络加一个上下文模块或者正在把 PyTorch 模型迁到 TF 部署这篇应该能帮你省不少时间。1. PSP模块的设计思路与核心原理1.1 为什么分割任务需要额外上下文普通卷积堆叠出来的感受野是“慢慢长”的。一个 3×3 卷积经过十几层之后理论上感受野能覆盖全图但实际有效感受野远没有理论值那么大中心区域占了绝大部分权重。这个现象在分割任务里非常致命远处一整片相似纹理可能被归成同一类单个小目标也可能因为周围环境相似而被吞掉。我常用一个生活类比来解释看一张几百人的合影如果只看每个人的脸你很难判断这些人站在哪里但如果你先看整体构图再看局部细节就能同时知道“这是个大场面”和“这一排人正在做什么”。PSP模块要解决的就是让网络同时拥有这两种“焦距”。它不靠增加卷积层数而是直接把特征图切成不同分辨率的网格让网络同时看到全图统计信息、区域统计信息和局部细节信息。这个思路和全局平均池化不一样。全局平均池化是把整张图压成一个点信息保留程度太低只靠多尺度卷积又会在计算量和参数上付出很大代价。金字塔池化正好折中网格越粗覆盖范围越大语义越全局网格越细越接近原始位置空间细节越完整。1.2 金字塔池化的设计逻辑PSP模块的输入通常是一张特征图形状记作 H×W×C。模块内部会并行做四路池化目标网格尺寸分别是 1×1、2×2、3×3、6×6。以 1×1 为例就是全局平均池化把整个特征图压缩成单个向量代表“整幅画面到底属于什么场景”。2×2 把画面分成四块每块告诉网络“左上角是天空、右下角是路面”。3×3 和 6×6 进一步细化保留更多空间对应关系。每一路池化之后特征图的尺寸会变小比如 24×24 的特征图经过 6×6 池化后变成 6×6。这时候如果直接拼接回 24×24尺寸对不上所以需要把每个池化结果做 1×1 卷积降维通常降到输入通道数的四分之一然后双线性插值上采样恢复成 H×W。最后把原始特征图和四路上采样结果按通道维度拼接再接一个 3×3 卷积做融合。维度上有一个容易忽略的点如果输入 C2048四路分支每路降到 512拼接后就是 2048 4×512 4096。这个通道数非常大所以最终融合卷积一般会再降到 512 或 256否则后续分类卷积计算量会失控。很多人在复现时只关注池化部分忽略了通道压缩导致显存直接爆掉。1.3 模块在常见分割网络中的位置PSP模块不是一个独立网络而是一个可插拔的积木。最经典的位置是放在主干网络之后、分类卷积之前。主干网络可以是 ResNet 这类通用分类网络最后几层去掉全局池化和全连接保留带空间分辨率的特征图PSP模块接在后面做多尺度上下文聚合最后再接一个 1×1 卷积输出每像素的类别概率。除了原版的 PSPNet这个模块也能塞进 U-Net 的 bottleneck、DeepLab 的 decoder 之前甚至放在实时分割网络的轻量主干后面。和 DeepLab 系列常用的 ASPP 相比PSP 更依赖池化操作而不是空洞卷积参数更少部署时对框架的算子要求也更低。正因为这种通用性我才特意在两个主流框架里各写了一份实现方便在不同项目里直接抄。2. PyTorch实现要点与代码2.1 选择 AdaptiveAvgPool2d 的原因PyTorch 里实现 PSP 模块最省心的地方就是torch.nn.functional.adaptive_avg_pool2d这个函数。它的含义很直接不管输入特征图是多大我都让你输出一个固定尺寸的网格。输入是 24×24目标 6×6它就自动把 24×24 分成 6 块区间输入换成 31×29它也能算不需要你手动去凑整。这个特性在实际工程里非常值钱。分割训练的输入尺寸经常不固定或者会在训练和推理时使用不同的分辨率。如果用普通的avg_pool2d必须自己算 kernel size 和 stride输入一变就得跟着改代码。用 AdaptiveAvgPool2d 就没有这个问题这也是 PyTorch 社区复现 PSP 模块几乎都这么写的原因。需要注意 AdaptiveAvgPool2d 的窗口划分和普通平均池化不完全一样它对边界块的处理是取整的每块覆盖的像素数可能略有差异。不过这种差异对语义分割这种任务来说基本没有影响因为后续还有 1×1 卷积和上采样特征图信息已经被高度抽象了。2.2 PyTorch 完整实现代码下面是我在项目里常用的 PSP 模块实现保留了足够多的注释方便对照维度变化。import torch import torch.nn as nn import torch.nn.functional as F class PyramidPoolingModule(nn.Module): def __init__(self, in_channels, out_channels512, bin_sizes(1, 2, 3, 6)): super().__init__() self.bin_sizes bin_sizes # 每个池化分支用 1x1 卷积压缩通道降到输入通道的 1/len(bin_sizes) inter_channels in_channels // len(bin_sizes) self.reduce_convs nn.ModuleList([ nn.Sequential( nn.Conv2d(in_channels, inter_channels, kernel_size1, biasFalse), nn.BatchNorm2d(inter_channels), nn.ReLU(inplaceTrue) ) for _ in bin_sizes ]) # 拼接后的融合卷积 self.fusion nn.Sequential( nn.Conv2d(in_channels len(bin_sizes) * inter_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): H, W x.shape[-2:] features [x] for size, conv in zip(self.bin_sizes, self.reduce_convs): pooled F.adaptive_avg_pool2d(x, output_sizesize) pooled conv(pooled) # 上采样回原始分辨率align_cornersFalse 与 TF 的 half_pixel_centersTrue 对齐 up F.interpolate(pooled, size(H, W), modebilinear, align_cornersFalse) features.append(up) concat torch.cat(features, dim1) return self.fusion(concat)这段代码的 forward 流程一目了然先取原始特征图作为第一路然后四路池化、卷积、上采样最后拼接融合。输入输出形状只差通道数空间尺寸保持不变所以后面可以直接接分类头。2.3 通道设计与网络组合通道设计上最常见的一组配置是主干输出 2048 通道四个分支每路降到 512拼接后 4096再通过融合卷积降到 512。这样 PSP 模块输出的通道数就和 DeepLab 的 ASPP 输出接近后续接 256 通道的分类头也比较顺。如果主干输出通道比较少比如只有 320我建议把inter_channels改成in_channels // len(bin_sizes)也就是每路 80拼接后 640融合卷积输出 256。这样参数不会爆炸MobileNet 这类轻量主干也能撑得住。在 PyTorch 里接入主干很简单定义一个包含 PPM 和分类卷积的分割头即可。class PSPHead(nn.Module): def __init__(self, in_channels, num_classes, mid_channels512): super().__init__() self.ppm PyramidPoolingModule(in_channels, mid_channels) self.cls nn.Conv2d(mid_channels, num_classes, kernel_size1) def forward(self, x): x self.ppm(x) return self.cls(x)实际训练时主干和分割头可以分别初始化。主干用分类数据集上预训练好的权重PSPHead 里的参数随机初始化然后一起做端到端微调。我试过先冻结主干只训练 PSPHead等损失降下来再解冻效果比一上来就全参数训练更稳。3. TensorFlow实现要点与代码3.1 TensorFlow 生态下先想清楚三件事写 TensorFlow 版之前有几件事必须提前想清楚否则代码写着写着就会卡住。第一TensorFlow 默认的数据格式是 channels_last也就是张量形状是 B×H×W×C。PyTorch 是 channels_first习惯上 B×C×H×W。如果你平时 PyTorch 写得多刚切到 TF 时会非常容易在维度操作上报错尤其是tf.concat和tf.image.resize的轴参数。第二TensorFlow 没有一个像 PyTorchAdaptiveAvgPool2d那样“一行代码输出任意尺寸”的常用高维接口。它的tf.nn.avg_pool需要你给定 kernel size 和 stride窗口大小得自己算。所以要么限制输入尺寸能被目标网格整除要么像我下面这样做一个简单的工程近似。第三建议直接用 Keras 函数式 API 搭模块而不是把池化和上采样都塞进自定义 Layer。因为tf.image.resize的目标尺寸可以是一个动态张量而 Keras 的UpSampling2D层要求静态size两者在动态输入下行为差很多。3.2 TensorFlow 完整实现代码下面是 TensorFlow 2.x 里的函数式实现。为了降低复现门槛我用一个自定义的adaptive_avg_pool_2d完成“任意目标尺寸”的近似先把输入缩放到目标尺寸的整数倍保证能被整除再做平均池化。这个方案在动态输入下很稳定。import tensorflow as tf from tensorflow.keras import layers def adaptive_avg_pool_2d(x, output_size): 近似 AdaptiveAvgPool2d。 先把特征图缩放到 output_size 的整数倍再做平均池化。 B, H, W, C tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2], tf.shape(x)[3] new_h output_size * tf.maximum(1, H // output_size) new_w output_size * tf.maximum(1, W // output_size) resized tf.image.resize(x, (new_h, new_w), methodbilinear, half_pixel_centersTrue) pool_h new_h // output_size pool_w new_w // output_size pooled tf.nn.avg_pool(resized, ksize[1, pool_h, pool_w, 1], strides[1, pool_h, pool_w, 1], paddingVALID) return pooled搭 PSP 模块时我把每一路的卷积和 BN 也放在函数里用函数式 API 串起来。def psp_module(input_tensor, out_channels512, bin_sizes(1, 2, 3, 6)): x input_tensor H tf.shape(x)[1] W tf.shape(x)[2] in_channels x.shape[-1] inter_channels in_channels // len(bin_sizes) branches [x] for size in bin_sizes: pooled adaptive_avg_pool_2d(x, size) pooled layers.Conv2D(inter_channels, 1, paddingsame, use_biasFalse)(pooled) pooled layers.BatchNormalization()(pooled) pooled layers.ReLU()(pooled) up tf.image.resize(pooled, (H, W), methodbilinear, half_pixel_centersTrue) branches.append(up) concat layers.Concatenate(axis-1)(branches) out layers.Conv2D(out_channels, 3, paddingsame, use_biasFalse)(concat) out layers.BatchNormalization()(out) out layers.ReLU()(out) return out这段代码和 PyTorch 版逻辑完全一致。注意psp_module只是一个搭模型的辅助函数只能在构建模型时调用一次不要在训练循环里反复调用否则每次都会创建一批新的卷积层。3.3 与预训练 backbone 集成接主干时我用 Keras 函数式 API 做一个完整的分割模型。主干可以选择带空洞卷积的版本保证输出分辨率不至于太小。backbone create_backbone() # 输出形状 B x H/16 x W/16 x C基于通用分类数据集预训练 ppm_out psp_module(backbone.output, out_channels256) head layers.Conv2D(num_classes, 1, nameseg_head)(ppm_out) model tf.keras.Model(backbone.input, head)BN 层在 TF 里由 Keras 自动管理训练时用model.fit会更新均值和方差推理时自动切到使用累计统计量。这一点比 PyTorch 手动model.train()/model.eval()稍微省心一点但有一个陷阱如果先冻结主干BN 依然会用累积均值做推理不会因为主干trainableFalse就停止工作。反而是在微调阶段如果主干 BN 被设置为可训练而分支 BN 也保持默认状态容易出现细微差异。我的习惯是微调阶段主干的 BN 保持trainableFalse只训练 PSP 模块里的 BN。4. 实操过程把小模块搭成分割头4.1 输入尺寸与数据准备在我自己跑过的项目里PyTorch 版对输入尺寸几乎零要求任意短边都能跑。但 TensorFlow 版为了配合我上面写的简化池化最好把训练尺寸定在一个能被 2 和 3 同时整除的数值上。以输出 stride 为 16 的 backbone 为例输入 384×384 时特征图是 24×24能被 1、2、3、6 都整除输入 480×480 时是 30×30同样能整除。这两个尺寸实测下来效果都不错。数据增强方面我一般用随机缩放 随机裁剪 左右翻转。裁剪尺寸会直接决定特征图是否能被池化网格整除所以如果你换了训练尺寸建议先打印一次特征图形状确认。不要想当然这个问题在 TF 版里真的很容易踩。4.2 训练配置与损失训练配置上主干一般加载分类预训练权重PSP模块随机初始化。损失函数用普通的交叉熵就能跑出不错的效果如果类别不均衡可以再加一个带难例挖掘的交叉熵。PSPNet 还有一个经典习惯在主干的中间层加一个辅助监督辅助损失权重通常在 0.4 左右。这个辅助头不是必须的但我实测加上之后早期训练更稳最终指标会高一点点。优化器我习惯用带动量的 SGD初始学习率 0.01配合 poly 学习率衰减策略。在很多分割数据集上这个组合比 Adam 更稳。如果是小数据集或者只想验证模块能不能跑通用 Adam 也可以但学习率要调低一些比如 0.001否则 BN 很容易在前期震荡。4.3 推理与后处理推理阶段模型输出是原始输入尺寸的 1/16需要双线性插值上采样回原图再做 softmax 和 argmax。多尺度推理是个便宜又有效的涨点手段常见做法是把输入缩放成 0.75、1.0、1.25 三个尺度分别预测再把 softmax 概率平均。加了 PSP 模块之后多尺度的收益依然存在因为全局上下文对不同尺度的响应会有差异。滑动窗口推理对 PSP 模块来说反而没那么合适。原因是窗口变小之后1×1 分支只在窗口内做全局池化会丢失整个图像的上下文信息。如果输入图特别大我更建议把整图等比缩放到一个合适尺寸做一次全局推理再用双线性插值把预测图放大回原图。4.4 一个小实验的直观观察我做过一个对比实验同一个 backbone同一个训练配置只差有没有 PSP 模块。加了之后最明显的变化是那些“大块物体边界处的零散误判”少了很多。比如一大片草地上有一小片裸露泥土原本只靠局部特征很容易把泥土和草地混在一起加了 PSP 之后网络能利用 1×1 和 2×2 分支的全局统计知道这块区域大概率属于草地中的异常从而更容易保留边界。另一个变化是训练曲线稳定得更快。PSP 模块相当于给网络提供了一个全局统计的“锚点”前几个 epoch 的 loss 下降速度比不加模块时明显更快。这个现象在轻量主干上尤其明显因为轻量主干本身的感受野不够池化分支正好补上了短板。5. 常见问题与排查技巧实录5.1 池化网格与特征图尺寸不整除我遇到最多的坑是 TensorFlow 端池化输出尺寸不对。举例来说特征图是 14×14目标池化到 6×6如果直接按 14//62 作为 stride 做平均池化输出会是 7×7不是 6×6。第一次跑的时候我甚至没留意到 shape 变了直到后面拼接时报维度错误才知道出了问题。解决办法有三个。第一训练尺寸固定成能被所有 bin_sizes 整除的值比如 384×384 搭配 stride 16这样最简单的窗口池化就够用。第二改用我上面写的adaptive_avg_pool_2d近似方案先缩放再做池化。第三动态输入下追求严格一致就要用循环把每个输出格子对应的输入区间算出来再对每个区间分别取均值。这个第三种方案代码量会大不少我一般只在做论文复现时才写。5.2 双线性插值的对齐差异PyTorch 的F.interpolate和 TensorFlow 的tf.image.resize表面看起来都是双线性插值但实际上像素坐标系定义不同。PyTorch 里align_cornersFalse对应坐标中点对齐TensorFlow 里要设置half_pixel_centersTrue才能得到一样的插值结果。反过来PyTorch 的align_cornersTrue对应 TensorFlow 默认的half_pixel_centersFalse。PyTorchTensorFlowalign_cornersFalsehalf_pixel_centersTruealign_cornersTruehalf_pixel_centersFalse这个差异在对齐两个框架的输出时特别重要。我一开始在 TF 端漏了half_pixel_centersTrue跑出来的特征图总是比 PyTorch 版在边缘处差一点点虽然不影响最终分割指标但会让模型移植时非常难排查。5.3 BatchNorm 在 Keras 里的 training 开关在自定义 Layer 里手写 PSP 模块时BatchNorm 的training参数很容易被忽略。如果你用的是 Keras 函数式 APImodel.fit和model.predict会自动切换训练/推理状态问题不大。但如果你把模块封装成自定义 Layer 并在call里手动调用子层就一定要加上training参数并传给 BatchNormalization否则验证阶段会拿着训练批次的统计量做推理结果波动很大。另外还有一个迁移学习的老坑给主干设置trainableFalse时Keras 会冻结全部权重但如果 BatchNormalization 的trainableFalse它用的是 moving_mean 和 moving_variance而不是当前 batch 的统计量。这对迁移学习是合理的但如果你希望微调阶段更新 BN 的统计量就不能简单地把所有层都设为trainableFalse需要单独处理 BN 层。5.4 显存占用与轻量化PSP 模块四路并行池化看起来会占不少显存实际上因为每个分支都先压通道再上采样显存增长并不夸张。真正吃显存的是拼接后的 4096 通道特征图。如果显存吃紧我第一推荐减少inter_channels比如从in_channels // 4降到in_channels // 8拼接通道数会立刻变小。第二推荐去掉 6×6 分支因为它在计算量和参数量上都偏高但对轻量主干来说收益有限。移动端部署时我更倾向于只用 1×1、2×2、3×3 三个分支并且把 fusion 卷积从 3×3 改成 1×1。这样 PSP 模块基本就是四五个 1×1 卷积加几次池化非常轻量对算力敏感的芯片也很友好。6. 选型建议与后续扩展6.1 两个框架输出是否一致如果你只是想验证模块实现得对不对可以在两个框架里分别搭一个随机初始化的 PSP 模块喂同样的输入只比较输出 shape 和数值量级。由于卷积核初始化不同数值不可能完全一致但形状必须相同且每个位置的数值应该在同一个分布范围内。如果要精确对比需要先把 PyTorch 的权重转换到 TensorFlow注意卷积核的排布方式不一样。PyTorch 的 Conv2d 权重形状是 [out, in, kh, kw]TensorFlow 的 Conv2D kernel 是 [kh, kw, in, out]转换时要 transpose。我在实际项目里转换过加上插值对齐之后模块输出的最大绝对误差通常能压到 1e-3 以下。6.2 什么时候选 PyTorch什么时候选 TensorFlow使用场景推荐框架理由论文复现、快速改结构PyTorchAdaptiveAvgPool2d更省心调试直观已有 TensorFlow/Keras 生态TensorFlow与现有训练推理链路无缝衔接移动端 / TFLite 部署TensorFlow量化工具链更成熟转换路径短大量分布式训练两者皆可取决于团队现有基础设施从我自己的经验看如果是新项目从零开始我倾向 PyTorch因为调试体验好写 PSP 这种动态尺寸模块几乎不用考虑对齐问题。但如果项目最后要落到移动端或者 TF Serving那我会直接用 TensorFlow 写省去两套框架之间转换权重的麻烦。6.3 模块的变体与扩展方向PSP 模块的下一步玩法其实很多。最简单的变体是调整 bin_sizes比如去掉 6×6或者改成 1、2、4不同数据集上最优组合不一样。还有一种思路是把池化后的特征经过 1×1 卷积后加一个注意力权重让网络自己决定哪个尺度的上下文更重要。这相当于把简单的拼接变成了加权融合涨点空间更大。另外PSP 和 FPN 这类特征金字塔结构并不冲突。可以把 FPN 各层的输出分别过 PSP 模块再逐层融合这样每个尺度的特征都有全局上下文兜底。代价是计算量上去了但效果通常也比单层加 PSP 更稳。我一般先不加只有当单层模块提升不明显时才往多尺度方向扩展。这几年每次接到分割任务我第一反应不是去换整个网络结构而是先往 baseline 里塞一个金字塔池化模块。它不一定是涨点最猛的却是我用过最稳的上下文模块之一。如果你也在 PyTorch 和 TensorFlow 之间来回折腾希望这篇小结能帮你把两个框架里的坑一次性趟平。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
