GANgealing 笔记（稠密视觉对齐，通过GAN监督学习，不需要对齐的数据集）

GANgealing 笔记（稠密视觉对齐，通过GAN监督学习，不需要对齐的数据集）

《GANgealing**😗* GAN-Supervised Dense Visual Alignment》

主页：https://www.wpeebles.com/gangealing

论文：https://arxiv.org/abs/2112.05143

项目：https://github.com/wpeebles/gangealing

英文讲解：BV1gS4y1D7TD

朱俊彦团队最新论文：GAN监督学习给晃动猫狗加表情，丝滑又贴合

---

看这篇的理由是，要实现隐式的重建动漫人体，我们需要一种能处理遮挡的稠密视觉对齐（也就是align），来作为坐标编码。

该方法能够很好的产生 align ，但是实验显示，该方法对遮挡的抗性很弱。不知道在训练时对 STN 的输入做随机 mask 能否克服这一点。

在对比实验中提到了一个 RAFT ，这个看起来效果也还行，但是在展示中， RAFT 有时会在某一帧突然崩坏，然后无法恢复。不过也许 RAFT 还是有可借鉴的地方，有空去看看。

这篇应该会提到很多 align 领域的文章，可以试着顺着爬一下看看。

流程总结

简单来说就是：

用 StyleGAN2 生成样本对 $(\mathbf{x},\mathbf{y})$ ，其中 $\mathbf{x}$ 是一个随机采样的不对齐的图片，而样本 $\mathbf{y}$ 是 $\mathbf{x}$ 的对齐的版本。然后训练一个 Spatial Transformer 将 $\mathbf{x}$ 对齐到 $\mathbf{y}$ 。

这里的关键是，产生 $\mathbf{y}$ 的过程是可学习的。我认为这是该方法区别于一般的用 GAN 合成数据训练 图像翻译 模型的主要特征。

$\mathbf{y}$ 是通过操纵 $\mathbf{x}$ 的隐编码产生的，而产生 $\mathbf{y}$ 的过程是有可学习的参数的，会跟着 Spatial Transformer Network 一起学习，从而可以产生更对齐的 $\mathbf{y}$ 。

$\mathbf{x}$ 由随机采样的隐编码 $\mathbf{w}$ 生成。 $\mathbf{y}$ 由混合出的隐编码 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 生成，其中 $\mathbf{c}$ 是一个可学习的隐编码。

产生 $\mathbf{y}$ 的方法非常关键：

我们想尽量让 $\mathbf{x}$ 与 $\mathbf{y}$ 的画面上是同一个东西，也就是说，要保证 STN 是有可能将这两者 align 的。于是我们根据 StyleGAN2 的特性，提出了用 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 操作隐编码来生成 $\mathbf{y}$ 。

$\mathrm{mix}(\mathbf{c},\mathbf{w})$ 是在 ${\mathcal{W}}^{+}$ 空间中完成的，也就是说，生成器的前几层用 $\mathbf{c}$ ，后面的层用 $\mathbf{w}$ 。这样做是因为，训练好的 StyleGAN2 就是有这样的性质，前几层控制姿态，后面的层控制外观，可以用实验验证这一点。

该方法能 work 有以下几个关键点：

1. StyleGAN2 能在不对齐的数据集上收敛，且能生成不对齐的图片。这点很重要，我本来以为这是不可能的。
2. 存在一种方法，可以将 StyleGAN2 生成的不对齐的图片对齐，或者说，存在一种方法，可以将 StyleGAN2 的生成结果的姿态和外观解耦。
   事实上我们知道这种方法是存在的，而且我们也知道这个方法的具体流程（还通过端到端的消融实验确认了流程细节和超参），我们只是把物体的要 align 到的具体姿态作为可学习变量而已。
3. 我们设计的 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 操作确实可以实现姿态和外观的解耦。并且，事实上，由于我们使用同一个 $\mathbf{c}$ 来控制物体的姿态，所以经过训练后，STN 也会自然的将所有图片对齐到同一个姿态。
4. Spatial Transformer 只能进行简单的仿射变换，这其实是一个较强的约束，这会一定程度上约束生成 $\mathbf{y}$ 的流程，令其与 $\mathbf{x}$ 之间尽量只有姿态差别。这使得，即使我们的 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 方法有可能无法完美地解耦姿态和外观，整个方法仍然有可能收敛。
   1. 我们本来可以直接用 GAN inversion 得到 $\mathbf{w}$ ，然后用 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 产生 align 的图像，但是还是训练了一个 STN ，就是因为，训练 STN 的过程也会反过来优化 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ ，这样可以比原本的方法更好。（另一方面是，实际上 GAN inversion 不是一个很鲁棒的方案，而且在非人脸数据上目前还不够好。用合成数据训练一个 STN 能获得更好的泛化性）

摘要

我们提出了 GAN监督学习 ，一种框架，可以端到端地联合学习 判别模型 和 对应的GAN生成的训练数据。

受经典的 Congealing 方法的启发，我们的 GANgealing算法训练了一个 Spatial Transformer，将随机样本从 在未对齐数据上训练的GAN 映射到 常见的 联合学习的 目标模式。

（后略，大致意思就是我们这个贼溜）

方法

我们提出了 GAN监督学习。

从一个预训练的 GAN 中采样出样本对 $(\mathbf{x},\mathbf{y})$ ，其中 $\mathbf{x}$ 为一个随机采样的样本，而样本 $\mathbf{y}$ 通过一种可学习的方式操作 $\mathbf{x}$ 的隐编码得到。

这个样本对用于训练网络 $f_{\theta }:\mathbf{x}\to \mathbf{y}$ ，用重建误差即可训练。

在普通的监督学习中， $f_{\theta }$ 是在固定的样本对上学习的，但是在 GAN监督学习中， $f_{\theta }$ 和 $\mathbf{y}$ 会端到端的联合学习。而在测试阶段，我们可以在真实样本上评估 $f_{\theta }$ 。

稠密视觉对齐

在这里，我们展示了 GAN监督学习 可以应用于 Congealing [48] —— 一种经典的无监督对齐算法。

在本例中， $f_{\theta }$ 是一个 Spatial Transformer 网络。

我们将这个算法称为 GANgealing。

首先，我们用未对齐的图片训练一个生成器 G ，从一个隐编码 $\mathbf{w}\in {\mathbb{R}}^{512}$ 生成 图像 $\mathbf{x}$ ，也就是 $\mathbf{x}=G(\mathbf{w})$ 。

固定目标

现在，考虑一个固定的隐编码 $\mathbf{c}$ ，以及合成出的固定图片 $G(\mathbf{c})$ 。

在传统 Congealing 中的一个简单的想法就是，将 $G(\mathbf{c})$ 作为目标模式 $\mathbf{y}$ ，也就是说，我们训练一个 Spatial Transformer 网络 T ，把所有随机的未对齐的图像 $\mathbf{x}=G(\mathbf{w})$ 都 warp 到同一个目标图像 $\mathbf{y}=G(\mathbf{c})$ 。

使用同一个可学习的目标

由于 G 是可微的，我们可以优化 $\mathbf{c}$ ，从而学习我们希望 congeal 的目标。具体来说，我们可以使用以下损失来联合学习 $\mathbf{c}$ 和 T ：

$${\mathcal{L}}_{\text{align }}(T,\mathbf{c})=\ell (T(G(\mathbf{w})),G(\mathbf{c}))$$

其中 $\ell$ 为图片之间的某种距离函数。

通过最小化 ${\mathcal{L}}_{\text{align }}$ 来优化 $\mathbf{c}$ ，我们可以尽可能地令 T 的工作变得简单。如果当前的 $\mathbf{c}$ 所对应的图像对于大部分图像都是没法用 Spatial Transformer 达到的，那么 $\mathbf{c}$ 会在优化中变得更容易达到。

这种方法对于多样性有限的数据集来说是合理的，然而在存在显著的外观和姿态变化的情况下，让每个图像都对齐到完全相同的目标图像是不合理的。这会导致不够好的结果，从表3的消融实验能够量化地体现出来（第一行）：

为每个样本使用不同的目标

一个理想的情况是，我们能为每个样本都构建一个目标，保留 $G(\mathbf{w})$ 的外观，但是目标之间的姿势、朝向都差不多相同。

为了实现这一点，给定 $G(\mathbf{w})$ ，我们令 $\mathbf{w}$ 的一部分等于 $\mathbf{c}$ ，来产生 $\mathbf{w}$ 对应的目标。具体地说，我们令 $\mathrm{mix}(\mathbf{c},\mathbf{w})\in {\mathbb{R}}^{512}$ 指代一个隐编码，其第一部来自于 $\mathbf{c}$ ，其余部分来自于 $\mathbf{w}$ 。由此我们可以提出 GANgealing 的损失函数：

$${\mathcal{L}}_{\text{align }}(T,\mathbf{c})=\ell (T(\underset{原本的\mathbf{x}}{\underbrace{G(\mathbf{w})}}),\underset{原本的\mathbf{y}}{\underbrace{G(\mathrm{mix}(\mathbf{c},\mathbf{w}))}})$$

其中 $\ell$ 为感知损失 [38] （又是你，看来得去看看这篇了）。

本文中，我们使用 StyleGAN2 作为 G，原则上可以用任何 GAN，不过StyleGAN2自带一定的 风格-姿态 解耦能力，我们可以利用这一点来构建上述的每个样本的目标。具体地说，我们使用风格混合 [42] 来构造 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ ， $\mathbf{c}$ 用于生成器的前几层， $\mathbf{w}$ 用于生成器的后面的层（也就是说，我们是用 ${\mathcal{W}}^{+}$ 空间来混合 $\mathbf{c}$ 和 $\mathbf{w}$ 的）。

表3展示了对于 “在何处开始输入 $\mathbf{w}$ ” 的消融实验（文中称为 cut off point）

对 Spatial Transformer 的参数化

回忆一下，一个 Spatial Transformer Network (STN)，输入一张图片，回归并且应用一个反向采样网格 $\mathbf{g}\in {\mathbb{R}}^{H\times W\times 2}$ （这里是把 STN 对图片的变化看作了 warp）。

本文主要研究一个执行相似性变换的 T （旋转平移，等比缩放）。我们也研究了一个无约束的 T ，直接回归像素级光流。我们最后的 T 是将相似性 STN 组成无约束 STN ，我们发现这样做效果最好。（就是说网络做的是分段仿射变换是吧）

与之前的工作相比 [50,64] ，我们发现多阶段训练没用，所以我们用端到端训练。最后，我们的 STN 也能在测试时执行水平翻转，详见附录B.4 。

当使用无约束的 T 时，加上一个光滑约束会有益，即：

$${\mathcal{L}}_{\mathrm{TV}}(T)={\mathcal{L}}_{\text{Huber }}\left({\Delta }_x\mathbf{g}\right)+{\mathcal{L}}_{\text{Huber }}\left({\Delta }_y\mathbf{g}\right)$$

其中 ${\mathcal{L}}_{\text{Huber }}$表示 Huber loss （就是分段函数，内侧L2，外侧L1）， ${\Delta }_x$ 和 ${\Delta }_y$ 表示有限差分下 x 和 y 的偏导。我们还对光流做 L2 约束，鼓励光流不做变换：

$${\mathcal{L}}_{\mathrm{I}}(T)=\Vert \mathbf{g}{\Vert }_2^2$$

对 c 的参数化

在实践中，我们不直接将梯度反向传播到 $\mathbf{c}$ 上，而是将 $\mathbf{c}$ 表示为 $\mathcal{W}$ 空间上 top-N 个主方向的线性组合 [29,78] ：

$$\mathbf{c}=\overline{\mathbf{w}}+\sum _{i=1}^N{\alpha }_i{\mathbf{d}}_i$$

其中 $\overline{\mathbf{w}}$ 表示经验平均 $\mathbf{w}$ 向量， ${\mathbf{d}}_i$ 表示第 $i$ 个主方向， ${\alpha }_i$ 是可学习的系数。

做这种重参数化的动机是，StyleGAN的 $\mathcal{W}$ 空间具有很强的表达能力。因此，在没有约束的情况下直接优化 $\mathbf{c}$ ，可能会远离真实图像的流形，产生比较差的目标图像。减少 N 的数量可以保持 $\mathbf{c}$ 处于流形上，并且预防退化解。表3中有对 N 的消融实验。

最终的 GANgealing 损失如下：

$$\mathcal{L}(T,\mathbf{c})={\mathbb{E}}_{\mathbf{w}\sim \mathcal{W}}\left[{\mathcal{L}}_{\text{align}}(T,\mathbf{c})+{\lambda }_{\mathrm{TV}}{\mathcal{L}}_{\mathrm{TV}}(T)+{\lambda }_I{\mathcal{L}}_{\mathrm{I}}(T)\right]$$

多个 c 联合对齐和聚类

如上所述， GANgealing 可以处理高度多样性的数据集。有些数据集，比如 LSUN House ，具有极其多样性的姿态，无法用单一模式很好地表示。所以我们同时使用 $K$ 个 $\mathbf{c}$ ，相当于对训练样本进行 $K$ 聚类，也就得到 { c k } k = 1 K \left\{\mathbf{c}_{k}\right\}_{k=1}^{K} {ck​}k=1K​ 。

使用 hard-assignment 来分配当前样本的类别，并计算损失：

$${\mathcal{L}}_{\text{align }}^K(T,\mathbf{c})=\underset{k}{\min }{\mathcal{L}}_{\text{align }}\left(T_k,{\mathbf{c}}_k\right)$$

注意每个类别使用不同的 $T_k$ 来 warp 图像。

在测试阶段，我们直接通过 $k^{\ast }=\arg {\min }_k{\mathcal{L}}_{\text{align }}\left(T_k,{\mathbf{c}}_k\right)$ 选定当前样本所属的类别。

当输入真实图象时，我们需要选择一个类别。一个合理的想法是直接用 GAN inversion [8, 11, 95] 反求出 $\mathbf{w}$ 然后算 $k^{\ast }$ 。不过在非人脸数据集上，当前的 GAN inversion 仍然很难还很慢，尽管最近有了一定的突破 [3, 33] 。

所以我们直接训练了一个分类器（用合成的数据），这样就行了。

实验

消融实验

第二列表示使用的感知损失，SSL 表示一个完全无监督的 VGG-16 [74] ，使用 SimCLR [13] 在 ImageNet-1k [17] 上预训练。我们发现 SSL 和 LPIPS 没有明显差别。

第三列表示 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 在第几层开始将 $\mathbf{c}$ 替换为 $\mathbf{w}$ 。

第四列表示令光流保持光滑的正则化系数。

第六列 PCK 表示正确检测的关键点所占百分比，越高越好。这是一个姿态估计里面的概念。

第一行表示使用不可学习的 $\mathbf{c}$ ，这展示了 GAN监督框架 联合学习判别模型和目标 的重要性。

第二行表示使用无约束的 $\mathbf{c}$ 进行学习，这展示了对 $\mathbf{c}$ 进行重参数化的必要性。

第三第四行对比了 $\mathrm{mix}(\mathbf{c},\mathbf{w})$ 在不同的位置将 $\mathbf{c}$ 替换为 $\mathbf{w}$ 的效果。

第五第六行表示使用同一个可学习的目标图像进行训练，而不进行mix，这是受到传统 Congealing 启发的 baseline。注意到这非常不稳定，而且 PCK 较低。

第七第八行展示了对光流进行光滑性约束的效果。