PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering 论文解读和感想

PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering 论文解读和感想

背景和动机

本文21年上旬发表于CVPR的关于无监督语义分割的文章。我个人不是专门做语义分割领域的，感觉无监督语义分割还是一块待开垦的沃土，做语义分割的小伙伴们可以尝试这个方向。首先本文针对之前无监督语义分割存在的两个问题：
1、分割时不同class的object像素点数量的不平衡问题
语义分割是像素级的分类问题，以往的无监督语义分割使用互信息的方式进行优化，然而由于不同class的object所包含的像素点个数不同（有时甚至差距很大），因此使用互信息会发生严重的样本不平和现象，由此导致的结果就是分割大大的偏向如天空，草地等占比大的objec（ leads to noisy representation as major classes）。
2、以往的无监督语义分割有一个重要的假设，即语义连贯性，然而在许多场景中，这种假设并不好，这会导致小object被忽略。

方法介绍

受到深度聚类方法（DeepCluster）的启发，本文提出了一种两阶段迭代的无监督语义分割方法：①首先通过特征提取网络$f_{\theta }(.)$将图像$x\in R^{3\times W\times H}$映射成特征图$f_{\theta }(x)\in R^{c\times w\times h}$，并将特征图看成是$wh$个$c$维的特征向量，之后用k-means对特征向量进行聚类，得到$K$个聚类中心和每个特征向量对应的label。当然以上过程并不是单张图像进行的，而是依batch进行的。②当对所有训练图像进行完上面步骤后，所有像素都有了对应的pseudo-label，然后用pseudo-label去监督特征提取网络$f_{\theta }(.)$的训练。
对于第①步k-means没有什么好说的，过程可以表示为
其中$[p]$代表第$p$个像素，${\mu }_{y_{ip}}$代表第$y$类特征向量的聚类中心。对于第二步，直觉是用一个分类器来进行，即

其中$g_w(.)$是对特征向量的分类器，是一个可训练的神经网络。然而这样做有一个非常大的问题，由于两个步骤都需要进行网络训练，就会造成训练难度大，使得最后收敛结果很差。为了避免这个问题，作者提出了一种无参数化的分类器，如下所示：

其中$d(.,.)$是计算两个特征向量的cosine距离。可以看到，在第①步得到聚类中心和pseudo-label后，作者直接拉近特征提取网络得到的特征向量到对应pseudo-label的聚类中心的距离。这里还有一点需要说一下，在代码实现时作者也实现的非常巧妙：作者构造了一个一层的$(c\times K)$的全连接网络，然后直接将得到的聚类中心当做网络参数，将得到的特征向量输入到这个网络中（我们暂且称之为$c(.)$），将得到的label和pseudo-label做交叉熵来更新$f_{\theta }(.)$。当然，在整个训练过程中$c(.)$的参数是固定的。
到这里，回到最开始的两个问题，第一个问题我们就可以给出解决方案了：在聚类结束后，统计每个类对应的像素点的数量，之后将其当做特征提取网络计算交叉熵时的权重，来抵消不同class样本不平衡带来的问题。
那么第二个问题是如何解决的呢，作者在本文提出了一种新的优化准则，即照片变换（ photometric transformations的不变性和几何变换的等变性。简单来说照片变换就是对输入图像做类似于颜色变换、明亮变换、高斯模糊等不改变像素点的语义信息（不改变得到的pseudo-label；几何变换则类似于随机翻转，随机剪切等。对图像做几何变换后再进行特征提取，等于对图像进行特征提取再做同样的几何变换。可以将以上两条准则归为以下监督损失：

其中


其中$P_i^{(1)}$和$P_i^2$是第$i$张图像两种不同的图像变换，而$G_i$则对应的是几何变换。通过这两种不同变换后得到的特征图其对应的pseudo-label是相同的。整个流程的伪代码如下