Labels4Free: Unsupervised Segmentation using StyleGAN论文解读

入选ICCV2021，官方page website，暂未开源
paper还是很有创意的，只是pipeline搭建起来很繁琐，还有很大的改进空间。

背景

styleGAN等生成对抗网络可以生成逼真的图像，生成器的中间特征已经有了前景和背景的特征学习能力，但目前没有方法研究让GAN生成图像的同时，还生成前景目标的mask。本篇paper正是为了解决该问题。
作者基于两种假设

• 前景和背景是独立的，因此独立生成然后组合，也能生成逼真的图像。
• 生成器的中间特征学习到了前背景特征分布，是可以利用起来，用于生成前景目标的mask的
  

Method

基于以上假设，作者将stylegan增加了一个Mask分支。同时整个框架分成3个部分。

• $G$，用于生成前景目标和mask，其中包含了一个预训练的stylegan和生成mask的$A$网络
• $G_{bg}$，用于生成背景图像，包含经过修改特征图的预训练的stylegan
• $D$，判别器，从头训练。
  

生成的图像通过以下公式得到：
$$G(z)\times A(z)+(1-A(z))\times G_{bg}(z^{\prime })$$

Alpha Network

生成mask的A网络，类似Unet的结构，从stylegan中间特征出发，设计了一个逐层上采样特征融合的轻量型模型。
在文中，作者提到的这一段，笔者并没有理解啥意思，如果有看懂的人，可以帮笔者解答一下。

生成背景的$G_{bg}$

作者刚开始是使用预训练的stylegan初始化的$G_{bg}$，但实验失败了，作者猜测这是因为stylegan已经学习了背景和前景的特征，训练出来的模型，生成的图像总是混杂了前景，无法用于仅生成背景。作者又在其他数据集（MIT places）上训练，发现判别器可以很轻松鉴别出不属于目标数据集(FFHQ等)的背景分布，导致生成的mask质量不佳。
因此，作者做了一种极具大胆创新的做法，他认为，stylegan已经学会了辨别前景背景。找到哪些channel对应了背景特征，然后把其他channel设置为0，就能把生成前景目标的功能去除。
那么如何找到这些channel呢？作者提出了一种基于计算梯度图的方式。

这个过程有点类似grad-cam，计算出哪些位置对目标函数的影响最大，则那些位置最有可能是前景位置，然后在统计每一个layer的gradient map，每个channel的梯度和。作者通过这种方式发现，第一个constant layer和w code注入的那个层，对应的梯度和是最大的。
作者提到，通过阈值去判断哪些channel需要设置0，会让部分背景依然还有前景目标，因此训练中是采用更加安全的措施，直接将指定的layer（constant layer和w code注入的那个层）设置为0。

训练过程

作者仅训练A和D，其他都是freeze的。没有对D使用path regularization，并且没有使用style mixing。
D是从头训练的，因此是一个弱鉴别器（博弈中弱于G），避免D在初期就很容易鉴别前景和背景。
还使用了一个regularization，但笔者没有看懂其中的含义。留待以后在研究。