【语义分割】ICCV21_Mining Contextual Information Beyond Image for Semantic Segmentation

文章目录

一、背景
二、动机
三、方法
- 3.1 整体结构
- 3.2 特征存储模块 Feature Memory Module
- 3.3 数据集层面的上下文聚合 Dataset-level Context Aggregation
- 3.4 表达一致性的学习 Representation Consistent Learning
- 3.5 Loss
四、效果

论文链接： https://arxiv.org/pdf/2108.11819.pdf
代码链接： https://github.com/CharlesPikachu/mcibi
代码合并到 sssegmentation： https://github.com/SegmentationBLWX/sssegmentation

一、背景

在语义分割中，探索不同类别的 “共现” 特征来建模的方法层出不穷，如 PSP/ASPP/OCR 等，但这些方法通常只在图像内建模，没有考虑到对单个图像以外的信息建模。但作者认为，深度学习模型是对一个数据集的整体数据学习，所以为了对像素进行更准确的分类，应该将其他图像中对应的相同类别的语义信息也加入学习中。

二、动机

所以为了缓解上面提到的问题，本文作者提出 “挖掘单个图像以外的上下文信息，来提升像素的特征表达”，如图 1 所示。

首先，在训练中，设置一个 “feature memory” 模块，来存储 dataset-level 的各个类别的历史输入
之后，预测当前图片的像素表达的概率分布，该概率分布是由真值的分布来监督训练的
最后，使用加权聚合的 dataset-level 的表达来增强每个像素表达，这里加权的全职是由相应的类概率分布决定的。

此外，为了在整个数据集层面进一步使得类内更聚合和类间更分散，作者设计了一个“表达一致性学习策略”，来使得分类头同时学习以下两者：

整个数据集层面：不同类别的 dataset-level 的表达
单个图像层面：pixel-level 的表达

三、方法

3.1 整体结构

首先，给定输入图像 III 经过 backbone 得到输出 R∈C×H8×W8R \in C \times \frac{H}{8} \times \frac{W}{8}R∈C×8H×8W，该输出矩阵中存储了图像的像素表达。

然后，使用如下方式来挖掘单个输入图像之外的上下文信息：

不同类别的 dataset-level 的表达存储在 feature memory module MMM 中
AbiA_{bi}Abi 是提出的 dataset-level 的上下文聚合机制
Cbi∈C×H8×W8C_{bi} \in C \times \frac{H}{8} \times \frac{W}{8}Cbi∈C×8H×8W 存储了来自于 MMM 的 dataset-level 的上下文聚合信息
H1H_1H1 是分类头，被用来预测像素表达的类别概率分布

为了将本文提出的方法嵌入现有的分割网络中，作者给所利用的网络设计了 self-existing context scheme，称为 AwiA_{wi}Awi，于是有：

CwiC_{wi}Cwi 存储了当前输入图像的上下文信息

之后，RRR 被增强：

FFF 是 transform function，被用于聚合原始表达 RRR，单个图像之外的上下文表达 CbiC_{bi}Cbi，图像内部的上下文表达 CwiC_{wi}Cwi

随后，使用 RaugR_{aug}Raug 来预测输入图像的每个像素的类别：

3.2 特征存储模块 Feature Memory Module

如图 2 所示，维度为 K×CK \times CK×C 的 feature memory module MMM 被用来存储不同类别的 dataset-level 的特征表达，MMM 的初始化采用随机初始化，在每次训练之后使用移动平均的方式来更新：

mmm：动量
ttt：是当前 iter
LLL：用来把 RRR 的维度处理的和 MMM 相同
mmm：使用多项式退火方法来确定
TTT：iteration 总数
ppp 和 m0m_0m0 都为 0.9

LLL 的实现：

首先，定义一个大小为 K×CK\times CK×C 的矩阵 R′R'R′，并使用 MMM 中的值来初始化该矩阵
RRR 被上采样，然后展成 HW×CHW\times CHW×C大小
然后，对于每个类别 ckc_kck，则有：
其中，GTGTGT of size HWHWHW 是真值
RckR_{ck}Rck of size Nck×CN_{ck} \times CNck×C 存储了类别 ckc_kck 的表达
NckN_{ck}Nck 是图像中真值为 ckc_kck 的像素个数
之后，计算 RckR_{ck}Rck 和 M[ck,∗]M_{[c_k, *]}M[ck,∗] 的相似矩阵 SckS_{ck}Sck of size NckN_{ck}Nck
最后，R′R'R′ 中的 ckc_kck 的表达被更新为：

3.3 数据集层面的上下文聚合 Dataset-level Context Aggregation

首先，预测一个大小为 K×H8×W8K \times \frac{H}{8} \times \frac{W}{8}K×8H×8W 的权重矩阵 WWW，来存储 RRR 中的特征表达的类别概率分布：

H1H_1H1：由两个 1x1 卷积和 softmax 函数组成

然后，计算粗糙的 dataset-level representation 矩阵 Cbi′C_{bi}'Cbi′：

Cbi′C_{bi}'Cbi′：大小为 HW64×K\frac{HW}{64}\times K64HW×K，存储了 dataset-level 的聚合表达
permute(W)permute(W)permute(W)：将 WWW 的维度转换为 HW64×K\frac{HW}{64}\times K64HW×K

由于 H1H_1H1 仅仅使用了 RRR 来预测 WWW，所以像素表达可能会类别错误，所以作者计算了 RRR 和 Cbi′C_{bi}'Cbi′ 的相关性，来获得位置置信权重来进一步 refine Cbi′C_{bi}'Cbi′，也就是：

首先，使用如下方式计算相关性 PPP：
- permutepermutepermute 用来将 RRR 的维度变为 HW64×K\frac{HW}{64}\times K64HW×K
然后，refine Cbi′C_{bi}'Cbi′
- permutepermutepermute 用来将 RRR 的维度变为 C×H8×W8C \times \frac{H}{8}\times \frac{W}{8}C×8H×8W

3.4 表达一致性的学习 Representation Consistent Learning

由于分割网络的目标是将整个数据集中的每个像素的特征表达影射到一个非线性空间中，但其训练时是通过 mini-batch 的方式来训练的，这种不一致的学习会导致网络缺失从整个数据集的角度来将同类拉近，将不同类分散开的能力，于是作者提出了 representation consistent learning strategy。

实现：

在训练时，使用 H2H_2H2 来预测 MMM 中的 dataset-level 表达的类别：

reshape：将 MMM reshape 到 K×C×1×1K\times C\times 1 \times 1K×C×1×1 的大小
H2H_2H2：由两个卷积层和一个 softmax 层构成
OMO^MOM：保存了 MMM 中的 dataset-level 的预测概率分布，MMM 中的每个表达都是对整个数据集的同一类别像素表达的整合
所以，在预测 OOO 和 OMO^MOM 的时候使用共享的分类头可以使得 H2H_2H2 能够 ① 提高图像内的每个像素的分类能力 ② 从整个数据集中学习到如何将类内拉近，将类间分散开。

3.5 Loss

该 Loss 是一个多任务 loss，涉及到 WWW、OMO^MOM、OOO

WWW 的 loss：

OMO^MOM 的 loss：

OOO 的 loss：

总体 loss：

α=0.4\alpha=0.4α=0.4
β=1\beta=1β=1

MMM 的值在反向传播中不更新

四、效果

学习到的特征的可视化：

在使用了本文的方法之后，学习到的特征的空间分布更明朗了，也就是同一个类别的像素表达更加集中，不同类别的像素特征表达更加分散。