附代码 ExFuse

ExFuse: Enhancing Feature Fusion for Semantic Segmentation论文解读

代码链接：https://github.com/lxtGH/fuse_seg_pytorch
参考链接：https://zhuanlan.zhihu.com/p/74551902

摘要：

在本文中，我们首先指出，由于在语义层次和空间分辨率上的差距，低级和高级特征的简单融合可能效果较差。我们发现，将语义信息引入低级特征，将高分辨率细节引入高级特征，对以后的融合更有效。基于此观察结果，我们提出了一个新的框架，名为ExFuse，以弥补低层次和高级特征之间的差距，从而显著提高了4.0%的分割效果。

背景：

低级特征和高级特征在本质上是互补的，其中低级特征的空间细节丰富，但缺乏语义信息，反之亦然。考虑一种极端情况，即“纯”低级特性（分辨率大）只编码低级概念，如点、线或边。直观地说，高级特征与这种“纯”低级特征的融合帮助很少，因为低级特征的噪声太大，无法提供足够的高分辨率语义指导。相反，如果低级特征包含更多的语义信息，例如，编码相对更清晰的语义边界，那么融合就会变得容易——通过将高级特征映射与边界对齐起来，就可以获得精细的分割结果。类似地，空间信息少的“纯”高级特征（分辨率小）不能充分利用低级特征；然而，由于嵌入了额外的高分辨率特征，高级特征可能有机会通过对齐到最近的低级边界来完善自己。
从经验上看，低级特征和高级特征之间的语义重叠和分辨率重叠对特征融合的有效性起着重要的作用。换句话说，可以通过在低级特征（大分辨率）中引入更多的语义概念或在高级特征（小分辨率）中嵌入更多的空间信息来增强特征融合。

提出了一个名为ExFuse的框架，它解决了这个差距从以下两个方面来看：

为了在低级特征中引入更多的语义信息，我们提出了层重新排列（Layer Rearrangement）、语义监督（Semantic Supervision）和语义嵌入分支（Semantic Embedding Branch）；
将更多的空间信息嵌入更多的高级特征，我们提出了两种新的方法：显式信道分辨率嵌入（Explicit Channel Resolution Embedding）和密集相邻预测（Densely Adjacent Prediction）。

网络整体结构：其中，使用GCN（Global Convolution Network）作为backbone。

训练顺序：

将图片通过一层卷积层
使用ResNet或者ResNeXt作为下采样网络，将特征映射下采样为四个不同分辨率的特征映射，即四个等级水平的映射。在预训练阶段，会加入语义监督SS在四个特征映射中，即加入辅助损失，一起进行训练，低级特征被迫编码更多的语义概念，预训练后，除去辅助损失，进行微调。
四个特征映射进入语义嵌入分支SEB，将低级特征与高级特征进行融合。res-5不进入。
接下来进入GCN。
最底下的最高级特征单独进入ECRE，使通道中含有分辨率的信息。上三层进入反卷积层。
最后进入密集相邻预测DAP，使模型可以预测邻近位置的结果。

GCN：

代码：

class _GlobalConvModule(nn.Module):def __init__(self, in_dim, out_dim, kernel_size):super(_GlobalConvModule, self).__init__()pad0 = int((kernel_size[0] - 1) / 2)pad1 = int((kernel_size[1] - 1) / 2)# kernel size had better be odd number so as to avoid alignment errorsuper(_GlobalConvModule, self).__init__()self.conv_l1 = nn.Conv2d(in_channels=in_dim, out_channels=out_dim, kernel_size=(kernel_size[0], 1),padding=(pad0, 0),bias = False)self.conv_l2 = nn.Conv2d(out_dim, out_dim, kernel_size=(1, kernel_size[1]),padding=(0, pad1),bias = False)self.conv_r1 = nn.Conv2d(in_dim, out_dim, kernel_size=(1, kernel_size[1]),padding=(0, pad1))self.conv_r2 = nn.Conv2d(out_dim, out_dim, kernel_size=(kernel_size[0], 1),padding=(pad0, 0))def forward(self, x):x_l = self.conv_l1(x)x_l = self.conv_l2(x_l)x_r = self.conv_r1(x)x_r = self.conv_r2(x_r)x = x_l + x_rreturn x

Layer Rearrangement（层重新排列）

为了使低级特征(上图res-2或res-3)“更接近”监督，一种直接的方法是在早期阶段安排更多的层，而不是后期。例如，ResNeXt101模型分别有{3、2、4、3、23、3}的构建块；我们重新安排分配到{8、8、9、8}，并调整通道的数量，以确保相同的整体计算复杂度。实验表明，尽管新设计的模型的ImageNet分类评分几乎没有变化，但其分割性能提高了0.8%。

Semantic Supervision（语义监督SS）

提出了另一种改进低级特征的方法，即语义监督(SS)，即将辅助监督直接分配到编码器网络的早期阶段（见上图）。为了在辅助分支中生成语义输出，低级特征被迫编码更多的语义概念，这将有助于以后的特征融合。我们的语义监督方法主要关注于提高低级特征的质量，而不是提升主干模型本身。
需要注意的是，这个模块并不是分割网络的一部分，而是在预训练的时候加在分类网络上的，总体分类损失等于所有辅助分支的加权求和。然后在预训练后，我们删除这些分支，并使用剩余的部分进行微调。。
结构如下：

Semantic Embedding Branch（语义嵌入分支SEB）

这个模块就是在把特征融合到decoder之前先将高低级别的特征进行融合。具体来说就是高级别的先经过卷积，再上采样，然后和低级别的进行逐像素相乘

Explicit Channel Resolution Embedding（显式通道分辨率嵌入ECRE）将分辨率信息嵌入到channels中

高级特征含有极少的分辨率，为了获得更多的细节，往往需要使用扩充策略，然而，扩充会带来极高的计算量，因此，我们将把分辨率信息嵌入到通道中。即通道中含有分辨率的信息

按照文中的说法，这个模块是加在第一个上采样环节的，对上采样的feature map加一个辅助loss，因为反卷积含有权重，参数是可学习的，辅助loss没办法嵌入模型中，因此使用Sub-pixel Upsample代替反卷积。但是这里是和谁做loss？这里只能理解为，该部分的上采样是直接变为原图大小，和label做loss，也就意味着这里的上采样并不是分割网络的主路的一部分，只是通过该模块对feature map施加影响。代码中使用nn.PixelShuffle完成。nn.PixelShuffle将（B，C, H, W）转换为（B，C/r^2, Hxr, Wxr）,将通道信息转换为分辨率信息。

Densely Adjacent Prediction（密度相邻预测DAP）：

传统的方法在最后预测分割结果是，每个像素点预测一个概率值，各个像素独立预测。而作者提出了一种方法，在预测某一位置像素的概率值时，参考周围3*3邻域的值，求平均得到。
空间定位 (i,j) 上的特征点主要负责相同位置的语义信息。为尽可能多地把空间信息编码进通道，本文提出一种全新的机制——密集邻域预测，可以预测邻近位置的结果，比如 (i-1,j+1) 。
原始GCN最后阶段为21通道（对应21个分类），因为是33邻域，一个像素需要参考9个像素的值，所以把21通道扩展为189（219）。
实现过程：先将通道信息转换为分辨率信息，即使用nn.PixelShuffle，再采用平均池化层获取。

代码：

import torch
from torch import nnfrom model.deeplab_resnet import ModelBuilderclass _GlobalConvModule(nn.Module):def __init__(self, in_dim, out_dim, kernel_size):super(_GlobalConvModule, self).__init__()pad0 = int((kernel_size[0] - 1) / 2)pad1 = int((kernel_size[1] - 1) / 2)# kernel size had better be odd number so as to avoid alignment errorsuper(_GlobalConvModule, self).__init__()self.conv_l1 = nn.Conv2d(in_channels=in_dim, out_channels=out_dim, kernel_size=(kernel_size[0], 1),padding=(pad0, 0),bias = False)self.conv_l2 = nn.Conv2d(out_dim, out_dim, kernel_size=(1, kernel_size[1]),padding=(0, pad1),bias = False)self.conv_r1 = nn.Conv2d(in_dim, out_dim, kernel_size=(1, kernel_size[1]),padding=(0, pad1))self.conv_r2 = nn.Conv2d(out_dim, out_dim, kernel_size=(kernel_size[0], 1),padding=(pad0, 0))def forward(self, x):x_l = self.conv_l1(x)x_l = self.conv_l2(x_l)x_r = self.conv_r1(x)x_r = self.conv_r2(x_r)x = x_l + x_rreturn xclass SEB(nn.Module):def __init__(self, in_channels, out_channels):super(SEB, self).__init__()self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1,padding=1)self.upsample = nn.Upsample(scale_factor=2, mode="bilinear")def forward(self, x):x1, x2 = xreturn x1 * self.upsample(self.conv(x2))class GCNFuse(nn.Module):def __init__(self, configer=None,kernel_size=7, dap_k=3):super(GCNFuse, self).__init__()self.num_classes =20num_classes = self.num_classesself.resnet_features = ModelBuilder().build_encoder("resnet101")self.layer0 = nn.Sequential(self.resnet_features.conv1, self.resnet_features.bn1,self.resnet_features.relu1, self.resnet_features.conv3,self.resnet_features.bn3, self.resnet_features.relu3)self.layer1 = nn.Sequential(self.resnet_features.maxpool, self.resnet_features.layer1)self.layer2 = self.resnet_features.layer2self.layer3 = self.resnet_features.layer3self.layer4 = self.resnet_features.layer4self.gcm1 = _GlobalConvModule(2048, num_classes * 4, (kernel_size, kernel_size))self.gcm2 = _GlobalConvModule(1024, num_classes, (kernel_size, kernel_size))self.gcm3 = _GlobalConvModule(512, num_classes * dap_k**2, (kernel_size, kernel_size))self.gcm4 = _GlobalConvModule(256, num_classes * dap_k**2, (kernel_size, kernel_size))self.deconv1 = nn.ConvTranspose2d(num_classes, num_classes * dap_k**2, kernel_size=4, stride=2, padding=1, bias=False)self.deconv2 = nn.ConvTranspose2d(num_classes, num_classes * dap_k**2, kernel_size=4, stride=2, padding=1, bias=False)self.deconv3 = nn.ConvTranspose2d(num_classes * dap_k**2, num_classes * dap_k**2, kernel_size=4, stride=2, padding=1, bias=False)self.deconv4 = nn.ConvTranspose2d(num_classes * dap_k**2, num_classes * dap_k**2, kernel_size=4, stride=2, padding=1, bias=False)self.deconv5 = nn.ConvTranspose2d(num_classes * dap_k**2, num_classes * dap_k**2, kernel_size=4, stride=2, padding=1, bias=False)self.ecre = nn.PixelShuffle(2)self.seb1 = SEB(2048, 1024)self.seb2 = SEB(3072, 512)self.seb3 = SEB(3584, 256)self.upsample2 = nn.Upsample(scale_factor=2, mode="bilinear")self.upsample4 = nn.Upsample(scale_factor=4, mode="bilinear")self.DAP = nn.Sequential(nn.PixelShuffle(dap_k),nn.AvgPool2d((dap_k,dap_k)))def forward(self, x):# suppose input = x , if x 512f0 = self.layer0(x)  # 256f1 = self.layer1(f0)  # 128print (f1.size())f2 = self.layer2(f1)  # 64print (f2.size())f3 = self.layer3(f2)  # 32print (f3.size())f4 = self.layer4(f3)  # 16print (f4.size())x = self.gcm1(f4)out1 = self.ecre(x)seb1 = self.seb1([f3, f4])gcn1 = self.gcm2(seb1)seb2 = self.seb2([f2, torch.cat([f3, self.upsample2(f4)], dim=1)])gcn2 = self.gcm3(seb2)seb3 = self.seb3([f1, torch.cat([f2, self.upsample2(f3), self.upsample4(f4)], dim=1)])gcn3 = self.gcm4(seb3)y = self.deconv2(gcn1 + out1)y = self.deconv3(gcn2 + y)y = self.deconv4(gcn3 + y)y = self.deconv5(y)y = self.DAP(y)return ydef freeze_bn(self):for m in self.modules():if isinstance(m, nn.BatchNorm2d):m.eval()if __name__ == '__main__':model = GCNFuse(20)model.freeze_bn()model.eval()image = torch.autograd.Variable(torch.randn(1, 3, 512, 512), volatile=True)res1= model(image)print (res1.size())