PSPNet

一.概述

Pyramid Scene Parsing Network（PSPNet）是CVPR2017上关于场景解析的文章，拿到了2016年ImageNet比赛中scene parsing任务的冠军，当然也常用来做语义分割。这篇文章出发点是在语义分割算法中引入更多的上下文信息（context information）， 这样能够避免许多误分割，PSPNet在FCN算法的基础上引入更多上下文信息是通过全局均值池化操作（global average pooling）和特征融合实现的，因此特征呈金字塔结构，这也是论文名叫pyramid的原因。PSPNet算法是目前应用比较广泛的语义分割算法之一，该算法在PASCAL VOC2012测试集上的mIOU是82.6%。其主要贡献如下：

• 在基于FCN的像素预测框架中，我们提出了一种金字塔场景解析网络
• 我们为深层的ResNet开发了一种有效的优化策略，基于深度监督的loss
• 我们为最先进的场景解析和语义分割构建了一个实用的系统，其中包含了所有关键的实现细节

FCN baseline存在的问题：

          

• Mismatched Relationship：上下文关系对理解复杂场景普遍十分重要。某种物体的出现往往具有一些固定的视觉模式。例如，一架飞机很可能在跑道上或者在空中飞行，而不是在道路上。例如图2的第一行，FCN根据外观将黄色方框中的船预测成了一辆“车”。但大家都知道，河水上基本不会出现车。缺乏收集上下文信息的能力增加了错误分类的可能性。
• Confusion Categories： ADE20K数据集中有许多类标签对十分难以区分。比如field和earth；mountain和hill；wall, house, building和skyscraper。这几对物体的外表十分相似。在图2的第二排，FCN将框中的对象一部分预测为了摩天大楼，一部分预测为了建筑。我们并不期望出现这种情况，整个对象应该是摩天大楼或建筑，但不应该是两者各有一点。此问题可以通过类别之间的关系来解决。
• Inconspicuous Classes：场景中往往包含任意大小的对象/东西。一些小尺寸的东西，如街灯和广告牌往往很难找到，但有时它们可能很重要。相反，大型的对象可能会超过FCN的感受野，从而导致不连续的预测。如图2第三行所示，枕头的外观与床单相似。不考虑全局场景中的类别的话可能会无法解析枕头。为了提高对于非常小或大的物体的性能，我们应该更加注意包含不显眼东西的子区域。

综上所述，许多错误都部分或全部与上下文关系和不同感受野的全局信息相关。因此，一个具有合适全局场景级先验的深层网络可以大大提高场景解析的性能。

二.详解

金字塔池化模块可以融合四种不同金字塔尺度的特征

• 红色突出显示的是最粗糙级别的单个全局池化bin输出
• 下面的金字塔分支将特征映射划分为不同的子区域，并形成针对不同位置的集合表示
• 金字塔池化模块中不同级别的输出包含不同大小的特征映射
• 为了维护全局特性的权重，如果金字塔共有N个级别，则在每个级别后使用1×1卷积，将对应级别的通道数量降为原本的1/N
• 然后通过双线性插值直接对低维特征图进行上采样，得到与原始特征映射相同尺寸的特征图
• 最后，将不同级别的特征concate起来，作为最终的金字塔池化全局特性

金字塔层级的数量和每一层的大小都可以进行调整

• 尺寸大小与输入金字塔池化层的特征映射的大小有关
• 该结构通过在几个stride中进行不同尺寸的池化来对不同的子区域实现抽样
• 因此，从表示性上来看，多个层级的核尺寸应该保持合理的差距

我们的金字塔池化模块是一个四层结构，bin大小为1×1，2×2，3×3和6×6。

                            

通过金字塔池化模块，我们提出的场景解析网络(PSPNet)如图3所示：

• 输入图像后，使用预训练的带空洞卷积ResNet提取特征图。最终的特征映射大小是输入图像的1/8，如图3(b)所示
• 在特征图上，我们使用(c)中的金字塔池化模块来收集上下文信息。使用4层金字塔结构，池化内核覆盖了图像的全部、一半和小部分。它们被融合为全局先验信息
• 在(c)的最后部分将之前的金字塔特征映射与原始特征映射concate起来
• 再进行卷积，生成(d)中的最终预测图

 总一个简单的总结即为：

• PSPNet为像素级场景解析提供了有效的全局上下文先验
• 金字塔池化模块可以收集具有层级的信息，比全局池化更有代表性
• 在计算量方面，我们的PSPNet并没有比原来的空洞卷积FCN网络有很大的增加
• 在端到端学习中，全局金字塔池化模块和局部FCN特征可以被同时训练
• ResNet起始的7x7的卷积替换为了三个3x3的卷积

这里还采用了一个附加的loss:

经过深度预先训练的网络能带来良好的性能，然而，网络深度的增加可能会带来额外的优化难度。

相反，我们提出了一个附加loss，可以监督初始结果，并用后面的最终loss来学习残差。因此，深层网络的优化被分解为两个问题，每一个都更容易解决。

                                               

• 除了使用Softmax loss来训练最终分类器的的主分支外，在第四阶段后再使用另一个分类器，即res4b22残差块
• 这两个loss同时传播，通过各自前面所有层
• 辅助loss有助于优化学习过程，主loss仍是主要的优化方向。增加权重，以平衡辅助loss

 在测试阶段，放弃辅助分支，只使用优化好的主分支来进行最终的预测。这种对基于ResNet的FCN进行深度监督的训练策略在不同的实验环境下是非常有用的，与预训练的ResNet模型也可以很好结合。体现了这种学习策略的普遍性。

三.代码分析

import torch
from torch import nn
from torch.nn import functional as F

import extractors


class PSPModule(nn.Module):
    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
        super().__init__()
        self.stages = []
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.stages] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)


class PSPUpsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.PReLU()
        )

    def forward(self, x):
        h, w = 2 * x.size(2), 2 * x.size(3)
        p = F.upsample(input=x, size=(h, w), mode='bilinear')
        return self.conv(p)


class PSPNet(nn.Module):
    def __init__(self, n_classes=18, sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet34',
                 pretrained=True):
        super().__init__()
        self.feats = getattr(extractors, backend)(pretrained)
        self.psp = PSPModule(psp_size, 1024, sizes)
        self.drop_1 = nn.Dropout2d(p=0.3)

        self.up_1 = PSPUpsample(1024, 256)
        self.up_2 = PSPUpsample(256, 64)
        self.up_3 = PSPUpsample(64, 64)

        self.drop_2 = nn.Dropout2d(p=0.15)
        self.final = nn.Sequential(
            nn.Conv2d(64, n_classes, kernel_size=1),
            nn.LogSoftmax()
        )

        self.classifier = nn.Sequential(
            nn.Linear(deep_features_size, 256),
            nn.ReLU(),
            nn.Linear(256, n_classes)
        )

    def forward(self, x):
        f, class_f = self.feats(x) 
        p = self.psp(f)
        p = self.drop_1(p)

        p = self.up_1(p)
        p = self.drop_2(p)

        p = self.up_2(p)
        p = self.drop_2(p)

        p = self.up_3(p)
        p = self.drop_2(p)

        auxiliary = F.adaptive_max_pool2d(input=class_f, output_size=(1, 1)).view(-1, class_f.size(1))

        return self.final(p), self.classifier(auxiliary)

四.实验结果