本文利用 VQ-VAE 的结构来学习姿态特征的编码，并用训练好的解码器和 codebook 来指导姿态估计模型的学习，把姿态估计任务变成了一个预测 codebook 中特征的分类问题。
下面是姿态估计模型比较通用的一个结构，目前的算法大都脱离不出这个范畴：
不知道大家有没有思考过这样一个问题：我们训练的姿态估计模型，算一个Encoder、Decoder、还是Encoder-Decoder结构？
为了简单起见，我们以 Topdown 算法来讨论。不难想到，至少有两种视角是说得通的：
-
整个模型（Backbone、Neck、Head）就是一个 Encoder，完成图片到位置信息的编码，位置信息的编码形式可以是具体的坐标，也可以是 Heatmap 或者离散的 bin。
-
Backbone 作为一个共识里的特征提取器，是 Encoder，将图片编码到了某个特征空间，然后 Head 部分是在对这份特征进行解码，解码成人类可以理解的位置信息，因此是 Encoder-Decoder 结构。
以上两种视角最大的不同在于，一种编码空间是人类定义的、可以理解的，另一种是模型自己学习的、人类不好理解或解释的。现阶段我们的知识并不能给出明确的判断，来证明或证伪模型内部的工作原理属于以上哪一种，亦或者其实都不属于，但大量的研究进展都证实，增加模型中可以解释、定义明确的部分、对模型学习的特征空间做出更多限定，对于我们优化算法、训练模型是有巨大帮助的。
在这篇文章中，我们来一起学习一下 CVPR 2023 的论文《Human Pose as Compositional Tokens》，文章利用 VQ-VAE 的结构来学习姿态特征的编码，并用训练好的解码器和 codebook 来指导姿态估计模型的学习，把姿态估计任务变成了一个预测 codebook 中特征的分类问题。
论文链接：https://arxiv.org/pdf/2303.11638.pdf
开源地址：https://github.com/Gengzigang/PCT/tree/main
在进入 PCT 论文的核心思想之前需要再对基础知识做一点铺垫，简单介绍一下 VQ-VAE。
其实近年的姿态估计研究中，利用生成模型的工作越来越多，比如我在《(论文笔记及思考：Human Pose Regression with Residual Log-likelihood Estimation(ICCV 2021 Oral))[https://zhuanlan.zhihu.com/p/395521994] 》中介绍过的 RLE ，使用 Normalizing Flows 中的 RealNVP 来建模数据的真实分布从而提升回归算法的精度；又比如扩散模型火了之后衍生出的好几篇 DiffPose 同名工作；我也读过一些用 GAN 或神经风格迁移来增加数据量、提升图片分辨率、数据增强的工作。这一次，在 PCT 中引入了 VQ-VAE。
不得不说，用到了生成模型的工作读起来都比较有意思，但唯一的问题在于数学公式比较多，理解起来会相对困难，直到现在也还是有同学反映 RLE 看不懂，我也在考虑再写一篇深入浅出 RLE ，来尝试尽量避开数学公式再解读一下（挖坑）。
AE
在讲 VQ-VAE 之前，我们可以先介绍一下最基础的 AutoEncoder(AE)。
基于 AE 的思想，很快也延伸出了 DAE 等工作，对输入图片进行人工的加噪处理，从而让模型可以更好地捕捉关键特征。
VAE
这里多说一句，为什么选择用高斯分布其实是有讲究的，高斯分布的好处是让特征都集中在0点附近，因此可以让特征空间变得稠密，这一来做生成任务的时候，随机取样取到的值也大概率是有意义的。
VQ-VAE
有意思的是，虽然 VQ-VAE 名字里带 VAE，但它其实更像一个 AE，因为训练好的 VQ-VAE 里 codebook 是固定的，这意味着你并不能像 VAE 一样随机取样送进 Decoder 来做生成，而必须通过某种方式预测 codebook 中的特征，组合之后才能给 Decoder 使用。
更细节的内容不在这里展开了，感兴趣的小伙伴可以自行学习。
Pose as Compositional Tokens (PCT)
基础知识铺垫完后，让我们进入 Human Pose as Compositional Tokens。顾名思义，这篇文章是想用 Token 的组合来表示人体姿态，什么叫“Token 的组合”呢？联系一下 VQ-VAE，自然是 codebook 里的特征了。
所以更通俗地理解，PCT 可以看成是用 Backbone 预测 M 个特征，然后把它们用 codebook 中与它们最近似的特征进行替换，然后扔给 Head 来预测坐标。
那么，很自然地我们就会问，codebook 里的特征从哪里呢？这里让我再一次请出 AE 的结构图：
那么，作为 AE 的上位替代，VQ-VAE 同样可以无缝衔接，我们只需要用姿态坐标来训练一个 VQ-VAE 网络，就可以得到 codebook，而这个 codebook 中的特征我们可以认为是各种不同姿态的构成元素，只要 codebook 足够大，那么我们一定能通过这些特征的组合来重建出原来的姿态。
由于输入是坐标不是图片，因此 Encoder 部分论文选择了 MLP-Mixer，这里可以简单地理解成一种不容易过拟合的大号的 MLP 了，其实理论上来说换成 Transformer 或者真的就用 MLP 也是没问题的，看你的任务难度来调整即可。
Decoder 部分也同样是 MLP-Mixer 的堆叠，用 M 个特征来重建 17x2 的姿态坐标。
这个网络是跟图片无关的，所以其实训练的时候其实理论上可以进行的数据增强会比带图片的要丰富很多，这一点做过 Pose Lifting 相关算法的同学应该知道我在说什么。
训练时的损失函数如下：
可以看到是 L1 loss 来监督重建的坐标值，第二项则是 VQ-VAE 中的 commitment loss，目的是鼓励 Encoder 的输出尽量接近 codebook 中的特征，防止 Encoder 预测的结果频繁在各个 codebook 特征之间跳动。
Class Head
训练好了姿态版本的 VQ-VAE，我们真正需要的其实只是 codebook 和 Decoder，原来的 Encoder 就可以扔掉了。接下来需要做的事是训练一个新的 Encoder，这个 Encoder 的能力是把输入图片也编码到 codebook 所在的特征空间，然后就可以快乐地查表和重建了。
但是要从头训一个 Backbone 来做这件事其实也挺麻烦的，作者想了一个更简单的办法，就是用在 Heatmap 方法上训好的 Backbone （MMPose 里一堆现成的）冻结住，在后面接一个轻量的 Class Head 做特征转换即可。
具体而言，Backbone 输出的特征图往往是维度比较高的，而 VQ-VAE 学出来的特征空间只需要 M 个特征，所以可以简单地通过 1x1 卷积降到 M 维，然后把二维的特征图拉直成一维，用全连接层变换到特定的特征维度。这里的做法跟 SimCC 中是比较接近的。
最后再通过一个分类层，对 M 个特征（形状为 MxN）进行分类，假设 codebook 中有 V 个特征，那么分类的结果就是 MxV 的 logits，然后像正常的分类问题一样在 [V] 维度取 softmax 就可以得到每个特征对应到 codebook 里的置信度。
在替换这一步，正常来做的话是对 MxV 的 logits 取 argmax 得到 codebook 里的索引，然后直接替换，但这样一来训练梯度就无法回传了，所以作者对这一步进行了软化，用 softmax 出来的结果乘上 codebook 里的特征矩阵：
也就是 (MxN) x (NxV) = (MxV) 的矩阵运算，从而使得梯度可以反向传播，网络也就可以端到端训练了。
可以结合作者开源的代码来理解以上计算过程。
# Head 部分
def forward(self, x, extra_x, joints=None, train=True):
"""Forward function."""
if self.stage_pct == "classifier":
batch_size = x[-1].shape[0]
cls_feat = self.conv_head[0](self.conv_trans(x[-1]))
cls_feat = cls_feat.flatten(2).transpose(2,1).flatten(1)
cls_feat = self.mixer_trans(cls_feat)
cls_feat = cls_feat.reshape(batch_size, self.token_num, -1)
for mixer_layer in self.mixer_head:
cls_feat = mixer_layer(cls_feat)
cls_feat = self.mixer_norm_layer(cls_feat)
cls_logits = self.cls_pred_layer(cls_feat)
encoding_scores = cls_logits.topk(1, dim=2)[0]
cls_logits = cls_logits.flatten(0,1)
cls_logits_softmax = cls_logits.clone().softmax(1)
else:
## 省略跟 class head 无关的代码 ##
## 省略跟 class head 无关的代码 ##
output_joints, cls_label, e_latent_loss = \
self.tokenizer(joints, joints_feat, cls_logits_softmax, train=train)
if train:
return cls_logits, output_joints, cls_label, e_latent_loss
else:
return output_joints, encoding_scores
# Tokenizer
def forward(self, joints, joints_feature, cls_logits, train=True):
"""Forward function. """
if train or self.stage_pct == "tokenizer":
## 省略跟 class head 无关的代码 ##
else:
bs = cls_logits.shape[0] // self.token_num
encoding_indices = None
if self.stage_pct == "classifier":
part_token_feat = torch.matmul(cls_logits, self.codebook)
else:
part_token_feat = torch.matmul(encodings, self.codebook)
if train and self.stage_pct == "tokenizer":
## 省略跟 class head 无关的代码 ##
else:
e_latent_loss = None
# Decoder of Tokenizer, Recover the joints.
part_token_feat = part_token_feat.view(bs, -1, self.token_dim)
part_token_feat = part_token_feat.transpose(2,1)
part_token_feat = self.decoder_token_mlp(part_token_feat).transpose(2,1)
decode_feat = self.decoder_start(part_token_feat)
for num_layer in self.decoder:
decode_feat = num_layer(decode_feat)
decode_feat = self.decoder_layer_norm(decode_feat)
recoverd_joints = self.recover_embed(decode_feat)
return recoverd_joints, encoding_indices, e_latent_loss
而我想说的是，相比于直接软化相乘，这里其实还可以有个更骚的操作：
quantize = cls_logits_softmax + (codebook - cls_logits_softmax).detach()
# 正向传播和往常一样
# 反向传播时，detach()这部分梯度为0，quantize和input的梯度相同
# 即实现将quantize复制给input
这样做可以做到数值上把 Encoder 预测的特征替换成 codebook 里的特征，但梯度上用预测的结果进行回传。好处是可以消除掉软化相乘的结果跟真实 codebook 之间的差异，因为软化出来的毕竟不是真的替换，特征还是会有细微的数值差异的，如果 Decoder 够强的话可能无所谓，但能消除差异肯定是更好的。有空也许可以验证一下这么做对于模型性能有没有帮助。
实验结果
从结果来看还是不错的，不过我注意到本文的实验都用的 Swin 这种非常强的 Backbone，不知道是不是因为冻结 Backbone 的做法对特征质量要求比较高，弱的 Backbone 可能效果不够。至于冻结 Backbone，嗯。。我闻到了经费紧张的味道。。whaosoft aiot http://143ai.com
本文的另一大卖点是对遮挡情况的鲁棒性，由于 VQ-VAE 是直接在姿态上进行训练的，所以理论上学到的 codebook 中的特征就可以任意组合出各种不同的姿态，而受图片本身质量的干扰会显著降低，在 OCHuman 和 CrowdPose 等数据集上的表现也佐证了这一点。
而 VQ-VAE 学出来的 codebook 特征由于含义非常明确，所以也很容易进行可视化验证，任意地调整输入 Decoder 的特征组合，就可以看到预测出来的姿态上的变化，并且这种变化是局部的：
结语
本文主要介绍了一种在姿态估计任务中引入 VQ-VAE 来约束姿态特征空间的方法。另外本文的官方代码是基于 MMPose 做的，也欢迎大家来试用 MMPose~