实时语义分割模型PIDNet（CVPR 2023）解析

paper：PIDNet: A Real-time Semantic Segmentation Network Inspired by PID Controllers

official implementation：GitHub - XuJiacong/PIDNet: This is the official repository for our recent work: PIDNet

third-party implementation： https://github.com/open-mmlab/mmsegmentation/tree/main/configs/pidnet

存在的问题

两分支结构在实时语义分割中展现出了其有效性，但高分辨率的细节信息和低分辨率的上下文信息的直接融合存在一个问题即细节特征很容易被周围的上下文特征给淹没。这种overshoot问题限制了现有two-branch分割模型精度的提升。

本文的创新点

本文将卷积神经网络CNN和PID (Proportional-Integral-Derivative)控制器联系起来，并表明一个两分支网络就相当于一个PI控制器，因此本质上存在类似的超调问题。为了解决这个问题，本文提出了一种新的三分支网络架构PIDNet，它包含三个分支分别用来解析细节信息、上下文信息和边缘信息，并采用边界注意力来指导detail branch和context branch的融合。截至发文时间，PIDNet是实时语义分割模型中的SOTA。

方法介绍

背景知识

一个PID控制器包含三个部分：比例proportional控制器、积分Integral控制器、微分Derivative控制器，如图3上所示。

PI控制器的实现如下所示 

比例控制器关注当前信号，积分控制器累积过去所有的信号，由于累积的惯性效应，当信号变成相反的时候，简单PI控制器的输出会发生超调现象。然后引入了微分控制器，当信号变小的时候，微分变量会变成负的，起到一个阻尼器的作用来减少超调。类似的，两分支网络Two-Branch Network (TBN)分别通过多个降采样和不降采样的卷积层来解析上下文和细节信息。考虑一个简单的一维例子，其中细节分支和上下文分支都包含3层卷积且不包含BN和ReLU，输出映射计算如下

其中 \(K_{i}^{D}=k_{31}k_{22}k_{13}+k_{31}k_{23}k_{12}+k_{32}k_{21}k_{13}+k_{32}k_{22}k_{12}+k_{32}k_{23}k{13}+k_{33}k_{21}k_{12}+k_{33}k_{22}k_{11}\)， \(K_{i}^{C}=k_{32}k_{22}k_{12}\)。这里 \(k_{mn}\) 表示第 \(m\) 层中卷积核的第 \(n\) 个值。由于 \(\left | k_{mn} \right | \) 大都分布在(0, 0.01)内（92% for DDRNet-23）并且以1为界，每项的系数会随着层数的增加呈指数下降。因此对于每个输入向量，更多的项意味着对最终输出有贡献的可能性更高。对于detail branch，\(I[i-1],I[i],I[i+1]\) 占所有项的70%还多，这意味这detail branch更关注局部信息。相反，在context branch中 \(I[i-1],I[i],I[i+1]\) 只占了所有项的不到26%，这表明context branch更关注周围的信息。图3下表明context branch对局部信息变化的敏感度低于detail branch。空间域中细节分支和上下文分支的这种行为类似于时域中的P (current) I (all previous) 控制器。

在PID控制器的z变换中用 \(e^{-j\omega}\) 替换 \(z^{-1}\)，如下

当输入频率 \(\omega\) 增加时，I控制器增益变小，D控制器增益变大，因此P , I, D控制器分别起到全通、低通滤波器、高通滤波器的作用。由于PI控制器更关注于输入信号的低频部分，且不能对信号的快速变化立刻做出反应，因此它本身就存在超调的问题。D控制器通过使输出对输入信号的变化更加敏感来减少超调。图3下展示了detail branch解析各种各样的语义信息尽管不准确，而context branch聚合低频上下文信息并且在语义上起到类似大型均值过滤器的作用。细节和语义信息的直接融合会导致一些细节特征的丢失，作者认为TBN在傅里叶领域等价于一个PI控制器。

PIDNet: A novel Three-branch Network

为了缓解超调问题，本文在TBN上增加了一个辅助微分分支（auxiliary derivative branch, ADB）来在空间上模拟PID控制器，并突出高频语义信息。每个对象内部像素的语义是一致的，只有在相邻对象的边界上才会不一致，所以语义的差只在边界上不为零，ADB的目标就是边界检测。为此，本文提出了一个新的三分支实时语义分割模型结构，Proportional-Integral-Derivative Network (PIDNet)，如图4所示

PIDNet包含三个互补的分支，比例P分支在高分辨率特征图中解析和保存细节信息，积分I分支在局部和全局聚合上下文信息来解析long-range dependencies，微分D分支提取高频特征来预测边缘区域。同时还采用cascaded residual block来构建硬件友好的骨干网络。此外P, I, D分支的深度分别设置为适中、深、浅来保证高效实现。

作者在第一个Pag module的输出处接了一个head来得到额外的语义损失 \(l_{0}\) 以便更好地优化整个网络。本文没有采用dice loss，而是采用加权交叉熵损失来处理边缘检测的不平衡问题。\(l_{2}\) 和 \(l_{3}\) 表示CE loss，具体 \(l_{3}\) 用的是boundary-awareness CE loss，利用边界head的输出来对齐语义分割和边缘检测任务，并增强Bag module的作用。BAS-Loss的计算如下

其中 \(t\) 是预先设定的阈值，\(b_{i},s_{i,c},\hat s_{i,c}\) 分别表示边界head的输出、分割的ground-truth和类别c第i个像素的预测结果。完整的损失函数如下

其中 \(\lambda_{0}=0.4,\lambda_{1}=20,\lambda_{2}=1,\lambda_{3}=1,t=0.8\)。

Pag: Learning High-level Semantics Selectively

在PIDNet中，I分支提供丰富准确的语义信息对细节解析和边缘检测至关重要，这两个分支的层数和通道数都相对较少。因此我们将I分支当做其它两个分支的后备，可以向其它两个分支提供需要的信息。和D分支直接相加的方式不同，本文提出了一个Pixel-attention-guided fusion module (Pag)如图5所示，让P分支有选择的从I分支学习有用的语义特征而不被淹没。

将P和I分支的特征图中对应像素的向量分别定义为 \(\vec v_{p}\) 和 \(\vec v_{i}\)，Sigmoid函数的输出可以表示为

其中 \(\sigma\) 表示这两个像素属于同一个对象的概率，如果 \(\sigma\) 的值大，我们更相信 \(\vec v_{i}\) 因为I分支在语义上是丰富而准确的，反之亦然。Pag的输出如下

PAPPM: Fast Aggregation of Contexts

为了更好的构建全局场景先验，PSPNet（PSPNet: Pyramid Scene Parsing Network）提出了pyramid pooling module (PPM)，它在卷积层之前拼接多尺度的pooling maps得到局部和全局的上下文表示。DDRNet（Deep Dual-resolution Network 原理与代码解析）中提出的Deep Aggregation PPM (DAPPM)进一步提高了上下文的embedding能力展现了优异的性能。但是DAPPM的计算过程就深度而言无法并行，这比较耗时。同时DAPPM的每个尺度都包含了太多的通道，可能超过了轻量模型的表示能力。因此，本文修改了DAPPM中的连接使其可以并行，如图6所示，并将每个尺度的通道数由128减少到96。

这个新的模块称为Parallel Aggregation PPM (PAPPM) 并且应用于PIDNet-M和PIDNet-S中以保证速度。在PIDNet-L中考虑到深度还是使用DAPPM，但减少了通道数来获得更快的速度。

Bag: Balancing the Details and Contexts

根据ADB提取的边缘特征，本文使用边界注意力机制来指导细节表示P和上下文表示I的融合。具体提出了一个Boundary-attention-guided fusion module (Bag) 如图7所示，分别用detail特征和context特征填充高频和低频区域。

context分支在语义上是准确，但丢失了太多的空间和几何细节，尤其是对于边界区域和小目标。由于detail分支保留了更多的空间细节，我们强迫模型在边界区域更信任detail分支，而在其它区域用context特征来填充。定义P, I, D特征图对应像素的向量分别为 \(\vec v_{p}, \vec v_{i}, \vec v_{d}\)，则Sigmoid、Bag、Light-Bag的输出分别如下

其中 \(f\) 表示卷积、BN、ReLU的组合。尽管把Bag中的3x3卷积替换为两个1x1卷积得到Light-Bag，两者的作用是相似的，即当 \(\sigma > 0.5\) 时模型更信任细节特征，否则更相信语义特征。

代码解析

这里以MMSegmentation中的实现为例，讲一下具体实现。输入shape为(16, 3, 480, 480)，网络结构是PIDNet-S。S、M、L的区别主要是backbone、ppm、head中的channel数量，num_stem_blocks和num_branch_blocks的个数，和ppm的类型。

backbone的实现在mmseg/models/backbones/pidnet.py中，forward实现如下

def forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor]]:
    """Forward function.

    Args:
        x (Tensor): Input tensor with shape (B, C, H, W).

    Returns:
        Tensor or tuple[Tensor]: If self.training is True, return
            tuple[Tensor], else return Tensor.
    """
    w_out = x.shape[-1] // 8
    h_out = x.shape[-2] // 8

    # stage 0-2
    x = self.stem(x)  # (16,64,60,60)

    x_i = self.relu(self.i_branch_layers[0](x))  # (16,128,30,30)
    x_p = self.p_branch_layers[0](x)  # (16,64,60,60)
    x_d = self.d_branch_layers[0](x)  # (16,32,60,60)

    comp_i = self.compression_1(x_i)  # (16,64,30,30)
    x_p = self.pag_1(x_p, comp_i)  # (16,64,60,60)
    diff_i = self.diff_1(x_i)  # (16,32,30,30)
    x_d += F.interpolate(
        diff_i,
        size=[h_out, w_out],
        mode='bilinear',
        align_corners=self.align_corners)
    if self.training:
        temp_p = x_p.clone()

    # stage 4
    x_i = self.relu(self.i_branch_layers[1](x_i))  # (16,256,15,15)
    x_p = self.p_branch_layers[1](self.relu(x_p))  # (16,64,60,60)
    x_d = self.d_branch_layers[1](self.relu(x_d))  # (16,64,60,60)

    comp_i = self.compression_2(x_i)  # (16,64,15,15)
    x_p = self.pag_2(x_p, comp_i)  # (16,64,60,60)
    diff_i = self.diff_2(x_i)  # (16,64,15,15)
    x_d += F.interpolate(
        diff_i,
        size=[h_out, w_out],
        mode='bilinear',
        align_corners=self.align_corners)
    if self.training:
        temp_d = x_d.clone()

    # stage 5
    x_i = self.i_branch_layers[2](x_i)  # (16,512,8,8)
    x_p = self.p_branch_layers[2](self.relu(x_p))  # (16,128,60,60)
    x_d = self.d_branch_layers[2](self.relu(x_d))  # (16,128,60,60)

    x_i = self.spp(x_i)  # (16,128,8,8)
    x_i = F.interpolate(
        x_i,
        size=[h_out, w_out],
        mode='bilinear',
        align_corners=self.align_corners)
    out = self.dfm(x_p, x_i, x_d)  # (16,128,60,60)
    return (temp_p, out, temp_d) if self.training else out

self.stem首先是两个3x3-s2的conv+BN+ReLU，然后是两个nn.Sequential层，每层包含2个BasicBlock，在第二层的第一个BasicBlock的第一个卷积层中进行下采样，通道数x2。最终self.stem的输出shape为(16, 64, 60, 60)。 

如图4所示，三分支从上到下分别为P, I, D，分别对应细节分支、语义分支、边缘分支。

首先是I分支，代码如下，其中channels=32，num_branch_blocks=3，BasicBlock和Bottleneck都来源于ResNet，具体介绍见ResNet。可以看出一共有3层，前两层的block是BasicBlock，最后一层是Bottleneck。stride=2，每一层都进行下采样，这样是为了增大感受野，提取更丰富的语义特征。

self.i_branch_layers = nn.ModuleList()
for i in range(3):
    self.i_branch_layers.append(
        self._make_layer(
            block=BasicBlock if i < 2 else Bottleneck,
            in_channels=channels * 2**(i + 1),
            channels=channels * 8 if i > 0 else channels * 4,
            num_blocks=num_branch_blocks if i < 2 else 2,
            stride=2))

然后是P分支，代码如下，其中num_stem_blocks=2。和I分支相比，这里默认stride=1不进行下采样，是为了保留更多的细节信息。同时I分支三层的block数分别为3、3、2，而这里为2、2、1，这是因为不下采样图像分辨率更大，减少block数从而减少计算量，和I分支的计算量保持平衡。

self.p_branch_layers = nn.ModuleList()
for i in range(3):
    self.p_branch_layers.append(
        self._make_layer(
            block=BasicBlock if i < 2 else Bottleneck,
            in_channels=channels * 2,
            channels=channels * 2,
            num_blocks=num_stem_blocks if i < 2 else 1))

 最后是D分支，代码如下，边缘分支和细节分支一样不进行下采样，同时相比于其它两个分支，边缘分支的层数更少，可能是因为边缘信息作为辅助分支不需要太强的学习能力，还能减少计算量。

self.d_branch_layers = nn.ModuleList([
    self._make_single_layer(BasicBlock, channels * 2, channels),
    self._make_layer(Bottleneck, channels, channels, 1)
])
self.d_branch_layers.append(
    self._make_layer(Bottleneck, channels * 2, channels * 2, 1))

Pag的实现如下，其中self.f_i和self.f_p是两个1x1-s1的卷积

def forward(self, x_p: Tensor, x_i: Tensor) -> Tensor:  # (16,64,60,60),(16,64,30,30)
    """Forward function.

    Args:
        x_p (Tensor): The featrue map from P branch.
        x_i (Tensor): The featrue map from I branch.

    Returns:
        Tensor: The feature map with pixel-attention-guided fusion.
    """
    if self.after_relu:  # False
        x_p = self.relu(x_p)
        x_i = self.relu(x_i)

    f_i = self.f_i(x_i)  # (16,32,30,30)
    f_i = F.interpolate(
        f_i,
        size=x_p.shape[2:],
        mode=self.upsample_mode,
        align_corners=False)  # (16,32,60,60)

    f_p = self.f_p(x_p)  # (16,32,60,60)

    if self.with_channel:  # False
        sigma = torch.sigmoid(self.up(f_p * f_i))
    else:
        sigma = torch.sigmoid(torch.sum(f_p * f_i, dim=1).unsqueeze(1))  # (16,32,60,60)->(16,60,60)->(16,1,60,60)->(16,1,60,60)

    x_i = F.interpolate(
        x_i,
        size=x_p.shape[2:],
        mode=self.upsample_mode,
        align_corners=False)

    out = sigma * x_i + (1 - sigma) * x_p
    return out

PAPPM的实现如下

def forward(self, inputs: Tensor):
    x_ = self.scales[0](inputs)
    feats = []
    for i in range(1, self.num_scales):
        feat_up = F.interpolate(
            self.scales[i](inputs),
            size=inputs.shape[2:],
            mode=self.unsample_mode,
            align_corners=False)
        feats.append(feat_up + x_)
    # [(16,96,8,8),(16,96,8,8),(16,96,8,8),(16,96,8,8)]
    scale_out = self.processes(torch.cat(feats, dim=1))  # (16,384,8,8)
    return self.compression(torch.cat([x_, scale_out],
                                      dim=1)) + self.shortcut(inputs)

下面是DAPPM的实现，可以看出区别主要在于DAPPM遍历每个尺度池化的输出，上采样后与前一个输出feats[i-1]相加后在经过一个卷积处理后再添加进列表feats中，因此当前尺度的结果依赖于上一步的输出，没法并行计算。而在PAPPM中，每个尺度的上采样输出feat_up都与第一个尺度池化的输出相加即x_，因此后面每个尺度的输出可以并行计算。

def forward(self, inputs: Tensor):  # (16,1024,8,8)
    feats = []
    feats.append(self.scales[0](inputs))

    for i in range(1, self.num_scales):
        feat_up = F.interpolate(
            self.scales[i](inputs),
            size=inputs.shape[2:],
            mode=self.unsample_mode)
        feats.append(self.processes[i - 1](feat_up + feats[i - 1]))
    # [(16,128,8,8),(16,128,8,8),(16,128,8,8),(16,128,8,8),(16,128,8,8)]

    return self.compression(torch.cat(feats,
                                      dim=1)) + self.shortcut(inputs)  # (16,256,8,8)

Bag和Light-Bag的实现非常简单，如下

# Bag
def forward(self, x_p: Tensor, x_i: Tensor, x_d: Tensor) -> Tensor:
    """Forward function.

    Args:
        x_p (Tensor): The featrue map from P branch.
        x_i (Tensor): The featrue map from I branch.
        x_d (Tensor): The featrue map from D branch.

    Returns:
        Tensor: The feature map with boundary-attention-guided fusion.
    """
    sigma = torch.sigmoid(x_d)
    return self.conv(sigma * x_p + (1 - sigma) * x_i)

# LightBag
def forward(self, x_p: Tensor, x_i: Tensor, x_d: Tensor) -> Tensor:
    """Forward function.
    Args:
        x_p (Tensor): The featrue map from P branch.
        x_i (Tensor): The featrue map from I branch.
        x_d (Tensor): The featrue map from D branch.

    Returns:
        Tensor: The feature map with light boundary-attention-guided
            fusion.
    """
    sigma = torch.sigmoid(x_d)

    f_p = self.f_p((1 - sigma) * x_i + x_p)
    f_i = self.f_i(x_i + sigma * x_p)

    return f_p + f_i