MoveNet复现，轻量级⼈体姿态估计模型的修炼之路
这天接到个新需求，需要实时检测⾃然场景下⽬标⼈体的关键点位置。
从算法⼯程师的⾓度来拆解下需求：
1、检测⼈体关键点位置，就是⼈体姿态估计任务嘛；
2、要实时，那么就是终端部署，服务端那传输延时就不考虑了。对了咱们硬件不⼤⾏，所以肯定是要轻量级模型的，分辨率也不能太⼤，剪枝量化蒸馏三件套也要做好打算；
3、“⾃然场景下的⽬标⼈体”，意思就是场景下可能有多⼈，但是我们只需要⼀个⽬标的关键点，要考虑如何区分（这点后⾯再展开说）。
02
⽅案探索
2.1 初见
之前的项⽬经验主要是⼈脸相关的分类、检测、分割以及OCR等，虽然没做过⼈体姿态估计，但是需求分析完，我的第⼀感觉是“应该很简单”，这份底⽓来⾃于之前做过的⼈脸关键点检测项⽬，当时的经验是，处理好数据、使⽤合适的Loss，随便⽤个剪枝后的轻量级模型都可以达到很⾼的精度，⽽且在嵌⼊式板⼦上能跑到10ms以下（都不需要PFLD之类的）。⽽⼈脸关键点有68点，⼈体也就17个点，不是简单多了。
于是直接拿出之前⼈脸关键点的代码，修改亿点点细节（调整模型输出、准备训练数据、修改DataLoader等），⽤少部分数据先跑跑看看。数据就⽤COCO和MPII就差不多了，后续可以⾃⼰从⽹上爬取⼀些来扩充，当然如果有⽬标场景的数据就更好了。
⼯欲善其事，必先抓只⼩⽩⿏。也就是需要⼀个评测指标，分类常⽤Accuracy、F1，检测常⽤mAP，⽽⼈体姿态估计⼜有所不同，请教了⾕哥（Google），⼀般是⽤PCK（PCKh）和OKS+mAP，简单来说，前者就是计算predict和label点的距离再经过head size normalize，后者就是⽤类似于⽬标检测中IOU的OKS计算相似度，然后再算AP。虽然前者在学术界⽬前已经很少使⽤（主要是刷榜刷得太⾼了），但是我们场景只需要单⼈，加上⼯程化简单，肯定是选择PCK，且考虑数据不⼀定都有head的标注，以及我们场景⽬标⼤⼩⽐较⼀致，最终决定使⽤⾃定义PCK，160分辨率下，距离⼩于5则算正确，其实PCK也算作⼀种acc，所以后⾯⽤acc指代。
第⼀次跑完，验证集acc达到了 0.81，感觉还⾏，于是开始往丹炉疯狂加料（加数据、加Data balance、微调loss、不同关键点loss设置不同weight、怀疑特征提取不够甚⾄把FastSCNN的backbone都拿来了、各种调参等等），⼀顿操作下来，终于有了⼀点提升：val-acc 达到了0.9869。
这么⾼，看来是成了，部署上板测试之前，习惯性在⾃⼰电脑上先跑跑，准备好可视化demo⼀看，傻眼了，除了head之类的⽐较准，⼿稍微⼀动就很容易检测不到。
⼀腔热⾎终于被冰冷的现实击溃，看来这个“应该很简单”的任务并不简单...⾃⼰挖的坑，跪着也要填完。稍微总结下失败的教训，然后准备⽼⽼实实从零开始。
⼈脸关键点检测直接回归坐标点的效果好，很⼤原因在于特征分布⽐较集中。⾸先是⼈脸⼤⼩⽐较⼀致，其次是⼈脸是刚体，各个关键点相对位置也⽐较固定，最后是除了侧脸，基本不会遮挡，⽬标特征也⽐较清晰。相较之下，⼈体⼤⼩分布就不⼀致，且各个关节活动范围很⼤，相对位置不固定，还存在各种位置遮挡和服装遮挡的情况，相⽐来说任务难了很多。同时由于关键点的分布范围太⼴，单纯⽤全连接层去回归，很有可能出现某⼀部分神经元训练的⽐较好，另⼀部分⽐较差，也就导致泛化能⼒很差。
2.2 反思
ferocity
古风漫画人物直接回归关键点的思路⾏不通，那就打开国门睁眼看世界——看看现在的主流⽅案。这篇2020年的综述就挺不错的：Deep Learning-Bad Human Po Estimation: A Survey。如同华⼭派著名的⽓宗和剑宗之争，⼈体姿态估计也有不同的派别，⽽且打得更热闹，还分了两个⽅向：
⽬标学习形式：直接回归点坐标 VS 回归Heatmap
整体检测流程：Top-Bottom VS Bottom-Up
带佳字的男孩名字原来早期的⼈体姿态估计就是直接回归关键点，看来我也只是在重复前辈的路⽽已。后来发现性能不好，于是改为回归heatmap的⽅式来训练。简单来说就是把原图缩⼩⼏倍得到特征图，这个特征图有关键点的位置值就是1、背景就是0，因为只设1太稀疏了，所以把1改为⼀个⾼斯核，这样也符合实际语义。
确定好学习形式后，再看看检测流程。上⾯我的⽅案其实就属于Top-Bottom，也就是先检测⼈，再裁剪出来对每个⼈去检测关键点；⽽Bottom-Up则相反，直接检测出所有的关键点，再依次组合成⼀个个⼈。这⾥和⽬标检测中的one-stage和two-stage其实有点神似，⼀般来说，Top-Bottom⽅式精度⾼，但是速度慢⼀点，⽽Bottom-Up速度快，但是精度差⼀点。
同时了解了⼀下⽬前⼀些优秀的开源⽅案（主要侧重⼀些相对轻量级模型，因此就不包含HRNet之流
了）：AlphaPo、OpenPo、Lightweight OpenPo、BlazePo、Simple Balines、Fast Human Po Estimation、Global Context for Convolutional Po Machines、Simple and Lightweight Human Po Estimation等，那就来吧。
2.3 尝试
2.3.1 浅尝辄⽌
对上⾯提到的各种潜在可⾏⽅案，分别进⾏测试。
修改上⽂的我的⼈脸关键点回归模型（Top-Bottom），把输出改为heatmap，相应调整数据处理和loss，训练完效果⼀般，val-acc85；
1、AlphaPo、OpenPo都属于⽐较⼤的模型，且GPU速度也只是勉强实时，我们是端侧+CPU+低性能板⼦，估计1秒以上了，所以暂时不考虑，同时发现有个Lightweight OpenPo（Bottom-Up）的项⽬，直接拿来上板跑了下，360ms，有点希望；
2、BlazePo，这⾥直接⽤了TNN的模型，同时发现他们还开源了腾讯光流实验室的⼀个姿态估计模型（Top-Bottom），⽐BlazePo 要好，于是直接测试了后者，上板200ms，且除了转换后的模型没有什么其他资料；
3、Simple Balines（Top-Bottom），⽤预训练模型训练，val-acc在0.95左右还⾏，就是换成轻量级⽹络后（mobilenet-v3），精度掉到0.9左右，且可视化很差；
4、Fast Human Po Estimation（Top-Bottom），相⽐于沙漏⽹络，stage砍半，channel砍半，就没其它改变了；使⽤了蒸
馏，loss是两部分，⼀部分是⽼师的预测结果，⼀部分是GT；效果还⾏，720ms略慢；
5、Global Context for Convolutional Po Machines（Top-Bottom），提供的预训练模型，可视化效果还⾏，400ms左右；
6、Simple and Lightweight Human Po Estimation（Top-Bottom），可视化效果还⾏，250ms；
2.3.2 集中⽕⼒
经过综合考虑，最后选择了Lightweight OpenPo来做优化。
⼀是它是Bottom-Up的模型，不需要额外训练⼈体检测器（对应额外的数据处理成本、模型训练成本以及推理耗时）；⼆是可视化来看，它的精度也是属于表现最好的⼏个。
要了解Lightweight OpenPo，那么就需要先提下OpenPo，它使⽤了VGG作为特征提取，加上两个header，分别输出关键点的heatmap和PAF的heatmap，通过多个stage迭代优化（每个stage的输⼊和输出都是这两个heatmap），最后通过MSE进⾏优化。关键点的heatmap好理解，PAF是什么呢？这是论⽂提出的part affinity fields，⼀般翻译为亲和⼒向量场，它代表了两个关键点之间的连接信息，个⼈觉得可以直观理解为关键点之间的⾻骼信息，实现的时候是求两个关键点之间的单位向量（代表了⽅向），然后给对应的heatmap 位置赋值。最后把多⼈检测问题转化为⼆分图匹配问题，并⽤匈⽛利算法求得相连关键点最佳匹配。
⽽Lightweight OpenPo作者测试发现原openpo的refinement stage的5个阶段，从第⼀个refinement stage1之后，正确率没有提升多少，但是运算量就加⼤了。并且还发现了refinement stage⽹络的关键点定位(heatmap)和关键点组合(pafs)两部分的两条分⽀运算，有不少操作是⼀样的，造成运算冗余。于是主要进⾏了以下优化：
怎么画年画1、新的⽹络设计。只采⽤initial stage+refinement stage两个阶段⽹络。（此处，我觉得到refinement stage2阶段的性价⽐⽐较北京到西双版纳
⾼，refinement stage1的正确率还是有点低，尽管运算量不⼤。但是refinement stage2再加了约19GFLOPs的同时正确率就提⾼了3%左右，性价⽐更好。）
2、更换轻量级backbone。⽤MobileNet替换，并⽤空洞卷积优化⽹络。并且进⾏了⼀系列对⽐实验，发现MoblieNet v1⽐MobileNet v2效果居然更好。
3、refinement stage中的关键点定位(heatmap)和关键点组合(pafs)部分的⽹络权值共享，减少操作计算量。最后两个卷积层才分出两个分⽀分别⽤来预测关键点的定位(heatmap)和关键点组合(pafs)。
4、原openpo的refinement stage中都是使⽤7x7的卷积核。作者却使⽤了三个连续的1x1, 3x3, 3x3的卷积核来代替7x7，并且最后的3x3是使⽤了空洞卷积, dilation=2。另外由于⽤3个卷积核代替原来的⼀个卷积核,⽹络层数变得很深，所以作者⼜加上了⼀个residual连接。
研究了下Lightweight OpenPo源码，继续开始往丹炉疯狂加料（改backbone、剪枝、加数据、调分辨率、蒸馏、各种调参等等），⼀顿操作下来，得到了模型：速度80ms，acc0.95（对应的蒸馏teacher acc0.98）
速度勉强能接受，精度还⾏，但是可视化测试稍微有点不准，可能跟数据也有⼀定关系。
03
峰回路转
3.1 转机
就在我犹豫要不要继续深⼊优化Lightweight OpenPo的时候，⽆意间刷到了这样⼀篇⽂章：《实时检测17个⼈体关键点，⾕歌SOTA 姿态检测模型，⼿机端也能运⾏》。⽂章介绍了⾕歌今年五⽉开源的⼀个轻量级⼈体姿态估计模型MoveNet，在tfjs有对应的api。
众⾥寻他千百度，蓦然回⾸，那⼈却在灯⽕阑珊处。这不就是上天为我准备的吗？赶紧⽤⽹页端测试了下，感觉还⾏，加上之前有转bodypix的tfjs模型到tf模型的经验，准备直接下模型本地跑跑，幸运的是，这次的MoveNet不仅有tfjs模型，还有tflite模型，这下⽅便多了。
通过官⽅博客的只⾔⽚语以及Netron可视化⽹络，对模型有了初步了解，官⽅宣称是基于CenterNet做的修改，但是修改了很多（The prediction scheme looly follows CenterNet, with notable change
s that improve both speed and accuracy）。具体来说，使⽤了mobilenetv2+fpn作为backbone，输出四个header，最后经过后处理输出⼀个最靠近图⽚中⼼的⼈的关键点。由此可知，它也是属于Bottom-Up的流派，同时通过关键点回归的范围进⾏加权，只输出最靠近中⼼的⼀个⼈的关键点，和我们的场景绝佳匹配，因为就Lightweight OpenPo⽽⾔，其余⼈的PAF信息其实也是多余的。
先测测速度吧，⾸先尝试上板安卓直接加载tflite模型，报错Didn't find op for builtin opcode 'RESIZE_BILINEAR' version '3'，之前⽤的sorflow:tensorflow-lite:2.0.0，尝试2.3.0后可以了。速度150ms左右⼀帧，连续跑在120ms左右。然后尝试⼀些推理框架（⽐如Tengine、NCNN、MNN等），这⾥我习惯使⽤TNN。结果tflite转tnn和pb转tnn均failed，debug看是不⽀持ARG_MAX算⼦，这是后处理⾥⾯的，于是⾃⼰简单写了个去掉后处理的mobilenetv2+fpn+4个header的模型转tnn，上板测试80ms左右。
考虑到Lightweight OpenPo已经优化了很多，继续压榨空间有限，且上⽂提到有冗余输出，⽽这个MoveNet直接就能跑到80多ms，且只输出⼀个⼈的关键点，还有⾕歌背书，因此痛下决⼼，准备抛弃Lightweight OpenPo采⽤MoveNet的⽅案，尝试复现它。同时我还准备了Plan B：如果复现不了，就直接导出tflite的backbone权重，然后新建⼀个pytorch模型加载权重，得到模型推理输出后通过
C++代码实现后处理。这样的缺点是没法⽤⾃⼰数据训练，⽆法迭代优化，且直接使⽤官⽅提供的预训练模型可视化来看精度也没有达到特别⾼。
3.2 复现
好吧，那就开始准备撸起袖⼦复现MoveNet了。
只根据⼀个模型结构要复现⼀个模型训练，说难不难，说简单也不简单。模型结构搭建照猫画虎即可，关键是数据如何构造、Loss如何选择没有任何信息，甚⾄还隐藏了⼤量细节，⽐如⾼斯核的构造、loss权重的设置、优化器等各种超参数设置。道阻且长，⾏则将⾄，只能⼀步⼀个脚印了。
事先说明，以下所有分析均是主观猜想，只是经过实验验证效果也不错，不代表⾕歌官⽅原始实现就⼀定是这样，仅供参考。
3.2.1 模型结构
要复现⼀个模型并不难，很多模型有论⽂、有官⽅代码或者有其他⼈复现的代码、哪怕是各种专栏博客的拆解分析，都能帮助很多，遗憾的是，MoveNet都没有，唯⼀能参考的只有两个东西：⾕歌官⽅博客的介绍，以及TFHub提供的训练好的模型。这也是这篇⽂章分享的初衷。
观察模型结构，主要分为三部分：backbone、header、后处理。
1、Backbone前⾯说过了很简单，即mobilenetv2+fpn，为了保证可控性和灵活性，我没有使⽤现成的
⼀些mobilenetv2的实现，⽽是⾃⼰⼀层层⾃⼰搭起来的（其实也没有太⼤必要，后续换backbone我也是直接在现成的模型代码上改的）；
2、Header的构建就更简单了，输⼊backbone的特征图，经过各⾃的⼏个卷积层，最后输出各⾃维度的特征图即可，难点在于要理解每⼀个header的含义，整体结构参考如下：
江苏教育网>舞台摄影经过不断的分析、猜想、测试，现在直接分享得到的结论吧。四个header我们分别命名
head_heatmap，head_center，head_reg，head_offt以便说明：
head_heatmap的维度是[N,K,H,W]，n是batchsize，训练时⾃⼰指定，预测时⼀般为1；K代表关键点数量，⽐如17；H、W就是对应的特征图了，这⾥输⼊是192x192，降采样4倍就是48x48；它所代表的意义就是当前图像上所有⼈的关键点的heatmap，注意是所有⼈的；
head_center的维度是[N,1,H,W]，这⾥的1代表的是当前图像上所有⼈的中⼼点的heatmap，你可以简单理解为关键点，因为只有⼀个，所以通道为1；它所代表的意义官⽅博客说是arithmetic mean，即每⼀个⼈的所有关键点的算术平均数，但是我实测这样效果并不好，我⾃⼰最终是取得所有关键点得最⼤外接矩形的中⼼点，当存在⼀些较远的关键点的时候，可能算术平均数可以很好的训练⼤部分距离近的点，但是对较远的点效果差点，⽽我⽐较关注⼿腕这种较远的点，按我这么取对每⼀个点学习起来差不多，这个就仁者见仁智者见智了，以⾃⼰场景实验结果为准；
head_reg的维度是[N,2K,H,W]，K个关键点，坐标⽤x,y表⽰，那么就有2K个数据，就是对应这⾥的2K通道；那么数据如何构造呢？
根据模型结构的拆解，就是在每个⼈的center坐标位置，按2K通道顺序依次赋值x1,y1,x2,y2,...，这⾥
的x、y代表的是同⼀个⼈的关键点相对于中⼼点的偏移值，原始Movenet⽤的是特征图48尺⼨下的绝对偏移值，实测换成相对值（即除以size48转换到0-1区间）也是可以的，可以稍微加快收敛，不过⼏乎没有区别；雪花英语
head_offt的维度是[N,2K,H,W]，通道意义⼀样都是对应K个关键点的坐标，只不过上⾯是回归偏移值，这⾥是offt，含义是我们模型降采样特征图可能存在量化误差，⽐如192分辨率下x=0和x=3映射到48分辨率的特征图时坐标都变为了0；同时还有回归的误差，这⾥⼀并训练了；
3、后处理相对来说⽐较⿇烦，因为有各种奇奇怪怪的操作，这⼀步只有结合Netron可视化⼀步步⼤胆猜测，结合官⽅博客认真分析，结构如下，可以看到不是常规的CNN结构。
