做语义分割的都知道，一张图里真正让你头疼的往往不是局部的边缘，而是“全局关系”。比如一个路灯杆，近看是细长条，模型完全能识别；可一旦路灯和远处的白色墙面连在一起，特征图上几个像素共享同一块感受野，网络就开始犯迷糊，要么把路灯并进墙里，要么把墙跟天空混成一片。这时候，让网络“退后一步看全景”比堆更多卷积更有效。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 卷积做融合。

维度上有一个容易忽略的点：如果输入 C=2048，四路分支每路降到 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_channels=512, 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_size=1, bias=False), nn.BatchNorm2d(inter_channels), nn.ReLU(inplace=True) ) for _ in bin_sizes ]) # 拼接后的融合卷积 self.fusion = nn.Sequential( nn.Conv2d(in_channels + len(bin_sizes) * inter_channels, out_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) 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_size=size) pooled = conv(pooled) # 上采样回原始分辨率，align_corners=False 与 TF 的 half_pixel_centers=True 对齐 up = F.interpolate(pooled, size=(H, W), mode='bilinear', align_corners=False) features.append(up) concat = torch.cat(features, dim=1) 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_channels=512): super().__init__() self.ppm = PyramidPoolingModule(in_channels, mid_channels) self.cls = nn.Conv2d(mid_channels, num_classes, kernel_size=1) 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), method='bilinear', half_pixel_centers=True) 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], padding='VALID') return pooled

搭 PSP 模块时，我把每一路的卷积和 BN 也放在函数里，用函数式 API 串起来。

def psp_module(input_tensor, out_channels=512, 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, padding='same', use_bias=False)(pooled) pooled = layers.BatchNormalization()(pooled) pooled = layers.ReLU()(pooled) up = tf.image.resize(pooled, (H, W), method='bilinear', half_pixel_centers=True) branches.append(up) concat = layers.Concatenate(axis=-1)(branches) out = layers.Conv2D(out_channels, 3, padding='same', use_bias=False)(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_channels=256) head = layers.Conv2D(num_classes, 1, name='seg_head')(ppm_out) model = tf.keras.Model(backbone.input, head)

BN 层在 TF 里由 Keras 自动管理，训练时用model.fit会更新均值和方差，推理时自动切到使用累计统计量。这一点比 PyTorch 手动model.train()/model.eval()稍微省心一点，但有一个陷阱：如果先冻结主干，BN 依然会用累积均值做推理，不会因为主干trainable=False就停止工作。反而是在微调阶段，如果主干 BN 被设置为可训练，而分支 BN 也保持默认状态，容易出现细微差异。我的习惯是，微调阶段主干的 BN 保持trainable=False，只训练 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//6=2 作为 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_corners=False对应坐标中点对齐，TensorFlow 里要设置half_pixel_centers=True才能得到一样的插值结果。反过来，PyTorch 的align_corners=True对应 TensorFlow 默认的half_pixel_centers=False。

这个差异在对齐两个框架的输出时特别重要。我一开始在 TF 端漏了half_pixel_centers=True，跑出来的特征图总是比 PyTorch 版在边缘处差一点点，虽然不影响最终分割指标，但会让模型移植时非常难排查。

5.3 BatchNorm 在 Keras 里的 training 开关

在自定义 Layer 里手写 PSP 模块时，BatchNorm 的training参数很容易被忽略。如果你用的是 Keras 函数式 API，model.fit和model.predict会自动切换训练/推理状态，问题不大。但如果你把模块封装成自定义 Layer 并在call里手动调用子层，就一定要加上training参数并传给 BatchNormalization，否则验证阶段会拿着训练批次的统计量做推理，结果波动很大。

另外还有一个迁移学习的老坑：给主干设置trainable=False时，Keras 会冻结全部权重，但如果 BatchNormalization 的trainable=False，它用的是 moving_mean 和 moving_variance，而不是当前 batch 的统计量。这对迁移学习是合理的，但如果你希望微调阶段更新 BN 的统计量，就不能简单地把所有层都设为trainable=False，需要单独处理 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

从我自己的经验看，如果是新项目从零开始，我倾向 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 之间来回折腾，希望这篇小结能帮你把两个框架里的坑一次性趟平。