PSPnet（Pyramid Scene Parsing Network）是一种深度卷积神经网络结构，用于场景解析和语义分割任务。该模型是由中科院自动化研究所于2016年提出，并在当年的ImageNet比赛中取得了最好的成绩。它最核心的思想就是利用金字塔池化模块（Pyramid Pooling Module）来捕捉不同尺度的上下文信息，以提高对图像语义的理解和分割的准确性。PSPnet显著的优点在于：
1.可以有效利用上下文信息：传统的语义分割模型如 FCN 等，缺乏依据上下文推断的能力，而 PSPNet 能够充分利用图像中的上下文关系，从而提高对复杂场景的理解和分割精度。
2.融合多尺度信息：金字塔池化操作可以同时捕捉到图像中的全局和局部信息，适应不同大小和形状的物体分割。不同尺度的特征融合有助于模型更好地理解图像的语义信息，对于小面积物体的分割效果也有显著提升，避免了因物体尺寸过小或过大而导致的分割不准确问题。
PSPnet的总体网络结构如下：
总体流程：
1.原始影像经过骨干网络提取特征，一般是一个卷积神经网络，常用的如ResNet50；
2.获取到特征图之后，将特征图输入到PPM模块（Pyramid Pooling Module），分别提取不同尺度的特征；
（1）特征图先经过一个池化操作，一般选择自适应池化AdaptiveAvgPool2d，根据不同的池化大小设置（论文中选择1，2，3，6），得到1x1，2x2，3x3，6x6的四个特征图；
（2）对这四个特征图分别进行卷积操作（包括批归一化和非线性激活操作），得到不同尺度的特征图；
（3）将这四个不同尺度的特征图进行上采样，得到和原始图像经过骨干网络后的特征图一样大小的四个特征图；
（4）将这四个上采样后的特征图以及原始图像经过骨干网络后的特征图，按通道维度进行拼接（cat操作），得到最终的特征图输出；
3.将金字塔池化模块输出的特征图再经过卷积，将通道数变成需要分类的类别数；
4.此时得到的最终输出大小和经过骨干网络提取后的特征图大小是一致的，所以需要再做一个上采样操作，将分割输出图恢复成和原始图像一样的大小，得到最终分割结果。
根据上面的分析，我们就可以很方便的用Pytorch把代码实现：
import torch
import torchvision.models as models
import torch.nn as nn
from torch.nn import functional as F
class PyramidPool(nn.Module):
def __init__(self, in_channels, out_channels, pool_size):
super(PyramidPool, self).__init__()
self.features = nn.Sequential(
nn.AdaptiveAvgPool2d(pool_size),
nn.Conv2d(in_channels, out_channels, 1, bias=True),
nn.BatchNorm2d(out_channels, momentum=0.95),
nn.ReLU(inplace=True)
)
def forward(self, x):
size = x.shape
#print('size-------------->',size)
output = F.upsample_bilinear(self.features(x), size[2:])
return output
def initialize_weights(*models):
for model in models:
for module in model.modules():
if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
nn.init.kaiming_normal(module.weight)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.BatchNorm2d):
module.weight.data.fill_(1)
module.bias.data.zero_()
class PSPNet(nn.Module):
def __init__(self, num_classes, pretrained=True):
super(PSPNet, self).__init__()
print("initializing model")
self.resnet = models.resnet50(pretrained=pretrained)
self.layer5a = PyramidPool(2048, 512, 1)
self.layer5b = PyramidPool(2048, 512, 2)
self.layer5c = PyramidPool(2048, 512, 3)
self.layer5d = PyramidPool(2048, 512, 6)
self.final = nn.Sequential(
nn.Conv2d(4096, 512, 3, padding=1, bias=False),
nn.BatchNorm2d(512, momentum=.95),
nn.ReLU(inplace=True),
nn.Dropout(.1),
nn.Conv2d(512, num_classes, 1),
)
initialize_weights(self.layer5a, self.layer5b, self.layer5c, self.layer5d, self.final)
def forward(self, x):
size = x.size()
x = self.resnet.conv1(x)
x = self.resnet.bn1(x)
x = self.resnet.relu(x)
x = self.resnet.layer1(x)
x = self.resnet.layer2(x)
x = self.resnet.layer3(x)
x = self.resnet.layer4(x)
x = self.final(torch.cat([
x,
self.layer5a(x),
self.layer5b(x),
self.layer5c(x),
self.layer5d(x),
], 1))
return F.upsample_bilinear(x, size[2:])
if __name__ == "__main__":
# 随机生成输入数据
rgb = torch.randn(16, 3, 512, 512)
# 定义网络
net = PSPNet(num_classes=8, pretrained=False)
# 前向传播
out_cls = net(rgb)
# 打印输出大小
print(out_cls.shape)
# 输出：
initializing model
torch.Size([16, 8, 512, 512])
可以看到，模型输出结果被分成了8类，并且大小和原始图像一致。这里，我用PSPnet模型对VOC2012语义分割数据集进行训练，学习率0.001，batch_size设置为16，训练200个epoch，总体精度达到0.959，各类精度如下：
2024-11-27 19:35:55,593 - __main__ - DEBUG - epoch199----> acc = 0.959041
2024-11-27 19:35:56,452 - __main__ - DEBUG - saveing ../checkpoints/pspnet.pt
2024-11-27 19:36:56,652 - __main__ - DEBUG - --------------------------------------
2024-11-27 19:36:56,652 - __main__ - DEBUG - |0|background|0.97|
2024-11-27 19:36:56,652 - __main__ - DEBUG - |1|aeroplane|0.52|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |2|bicycle|0.23|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |3|bird|0.65|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |4|boat|0.5|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |5|bottle|0.43|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |6|bus|0.55|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |7|car|0.62|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |8|cat|0.77|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |9|chair|0.54|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |10|cow|0.49|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |11|diningtable|0.56|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |12|dog|0.73|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |13|horse|0.47|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |14|motorbike|0.53|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |15|person|0.89|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |16|potted plant|0.36|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |17|sheep|0.45|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |18|sofa|0.63|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |19|train|0.61|
2024-11-27 19:36:56,653 - __main__ - DEBUG - |20|tv/monitor|0.52|
2024-11-27 19:36:56,653 - __main__ - DEBUG - --------------------------------------
可以看到，人的识别精度还是比较高的，达到了0.89，但是对自行车的识别率很低，才0.23，我们用测试数据集的数据进行预测，预测结果如下：
可以看到总体精度还不错，人的边界还是比较清晰的。
再用武汉大学的GID数据集进行训练（GID 数据集（Gaofen Image Dataset）由武汉大学夏桂松团队制作。它包含从中国60多个不同城市获得的150张高分辨率Gaofen-2（GF-2）图像（7200×6800）。这些图像覆盖了超过50,000平方公里的地理区域。 GID中的图像具有高的类内分集以及较低的类间可分离性。分为建成区、农用地、林地、草地、水系与未标记区域，共6类。），训练100个epoch，总体精度达到0.89，具体精度如下：
2024-11-29 13:30:56,346 - __main__ - DEBUG - epoch99----> acc = 0.887112
2024-11-29 13:30:57,644 - __main__ - DEBUG - saveing ../checkpoints/pspnet2.pt
2024-11-29 13:38:48,590 - __main__ - DEBUG - --------------------------------------
2024-11-29 13:38:48,591 - __main__ - DEBUG - |0|background|0.84|
2024-11-29 13:38:48,591 - __main__ - DEBUG - |1|building|0.94|
2024-11-29 13:38:48,591 - __main__ - DEBUG - |2|farmland|0.94|
2024-11-29 13:38:48,591 - __main__ - DEBUG - |3|tree|0.39|
2024-11-29 13:38:48,591 - __main__ - DEBUG - |4|grass|0.63|
2024-11-29 13:38:48,591 - __main__ - DEBUG - |5|water|0.85|
2024-11-29 13:38:48,591 - __main__ - DEBUG - --------------------------------------
test loss = 0.263143
accuracy = 0.890584
可以看到，相比之下，PSPnet的解译结果是最平滑的，并且河流也识别出来了，虽然有少量的误分类。
