复现、并改进open-mmlab的mmpose详细细节

1.配置环境

stage1：创建python环境

conda create --name openmmlab python=3.8 -y
conda activate openmmlab

stage2：安装pytorch（这里我是以torch1.10.0为例）

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113

stage3：安装MMCV，全部都可以加上清华源，-i https://pypi.tuna.tsinghua.edu.cn/simple

pip install openmim
pip install mmengine
pip install "mmcv>=2.0.1"
pip install "mmdet>=3.1.0"

但是，这里在安装mmcv的时候，容易出错，比如：ImportError: DLL load failed while importing _ext: 找不到指定的程序。
很大就是mmcv版本装的有问题。
我们定位到mmcv的安装文档

https://mmcv.readthedocs.io/en/latest/get_started/installation.html#install-mmcv

会有根据操作系统、cuda版本，以及torch的版本后，确定自己选择的mmcv版本，进而生成最终的安装指令。

stage4：安装MMPose：
方式1：

git clone https://github.com/open-mmlab/mmpose.git
cd mmpose
pip install -r requirements.txt
pip install -v -e .

方式2：

pip install "mmpose>=1.1.0"

2.数据处理

这里我们以 CrowdPose 数据集为例，下载链接为：

https://github.com/MVIG-SJTU/AlphaPose/blob/pytorch/doc/CrowdPose.md

这里我们任意选取configs/body_2d_keypoint文件中的文件，以：topdown_heatmap/crowdpose/td-hm_res50_8xb64-210e_crowdpose-256x192.py为例。
新建文件夹CrowdPose，并将下载下来的数据集文件夹存放进去：

我们首先更改选好的配置文件td-hm_res50_8xb64-210e_crowdpose-256x192.py中的data_root，路径就是刚才我们新建好的文件夹。

接着，还是同一份配置文件，我们更改标注文件，这里文件的默认为：'annotations/mmpose_crowdpose_train.json'，我们可以根据这个来修改我们下载下来的数据集文件夹名称，也可以将这里反过来改成我们自己的数据集文件夹名，都可以的。

3.训练

4.改进mmpose

MMPose中，所有与模型结构实现相关的代码都存放在 models目录下：

4.1 快速调试技巧

mmpose是基于pytorch开发的，所以在数据预处理以及前向传播的调用，都在底层环境的mmengine中：/mmengine/model/base_model/base_model.py。

我们这里可以打断点过来，可以看到，image的维度是[batch_szie, 3， 256， 192]；

随后，我们可以在我们选用的backbone，或者自己新加的backbone中，在前向传播forward函数中打上断点，会发现这里的x就是上述的image，与此同时，也能看到每一层网络结构后feature map的变化。

这里，给推荐另一个小技巧，因为每次从整个框架的train开始断点，耗时麻烦。
我们直接在backbone的py文件中，进行调试，这样就大大减少了我们改进或者调试模型的复杂度，我们只要添加如下代码，将config配置文件中模型的extra 字典复制过来，然后在当前py文件进行debug就省事省力了。

if __name__ == '__main__':extra = dict(stage1=dict(num_modules=1,num_branches=1,block='BOTTLENECK',num_blocks=(4,),num_channels=(64,)),stage2=dict(num_modules=1,num_branches=2,block='BASIC',num_blocks=(4, 4),num_channels=(32, 64)),stage3=dict(num_modules=4,num_branches=3,block='BASIC',num_blocks=(4, 4, 4),num_channels=(32, 64, 128)),stage4=dict(num_modules=3,num_branches=4,block='BASIC',num_blocks=(4, 4, 4, 4),num_channels=(32, 64, 128, 256)))model = HRNet(extra, in_channels=3)model.eval()x = torch.randn((2, 3, 256, 192))outputs = model(x)

4.2 快速定位

将backbone、neck、head各结构实例化的接口位于mmpose/models/pose_estimators/base.py中。

4.3 改进backbone

4.3.1 使用说明

如果希望实现一个新的backbone，你需要在目录 backbones下新建一个文件进行定义。

新建的骨干网络需要继承BaseBackbone 类，其他方面与你继承 nn.Module 来创建没有任何不同。

在完成骨干网络的实现后，你需要使用 MODELS来对其进行注册：

from mmpose.registry import MODELS
from .base_backbone import BaseBackbone
@MODELS.register_module()
class YourNewBackbone(BaseBackbone): 

最后，请记得在backbones/__init__.py 中导入你的新骨干网络。

4.3.2 改进案例

以复现mmpose配置文件：configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192.py为例，该配置文件的backbone对应的配置文件为：‘HRNet’，我们比如想替换：

https://github.com/xzz777/SCTNet

中的开源项目中的分割backbone。

4.3.2.1 复现mmpose原配置文件

为了确保替换的backbone能跑通，最好是跑一遍原始配置文件，这样清楚知道模型的backbone、neck、head各结构输出的维度，这样方便在新模型中进行调整。
我们首先复现官方提供的原配置文件：configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192.py。
配置文件中：input_size表示dataloder预处理后图片的size（也就是模型的输入图像大小），heatmap_size表示最后输出的key points热图的size，这里可以更改为指定的size，比如：640，相应的热图size相应缩放4倍。

backbone输入：[2, 3, 640, 640]，backbone的输出为：[batch_size, 32, 160, 160]。

Head表示输出头，in_channels为32，表示backbone的输出通道数（如果没有设置neck结构的话），out_channels为17，表示Head的输出通道数，结合上述的heatmap_size最终的输出为：[batch_size, 17, 160, 160]。

4.3.2.2 复现开源项目

我们想用：

https://github.com/xzz777/SCTNet

中的开源项目中的分割backbone，那复现一下看一下底层的feature map的变化对于我们快速改完模型是帮助很大的。
输入：[2, 3, 640, 640]，输出如下：

4.3.2.3 修改配置文件

我们参照4.3.1 使用说明新建好开源项目中的backbone后，我们对配置文件中的model字典文件进行更改：

# model settings
model = dict(type='TopdownPoseEstimator',data_preprocessor=dict(type='PoseDataPreprocessor',mean=[123.675, 116.28, 103.53],std=[58.395, 57.12, 57.375],bgr_to_rgb=True),# backbone=dict(#     type='HRNet',#     in_channels=3,#     extra=dict(#         stage1=dict(#             num_modules=1,#             num_branches=1,#             block='BOTTLENECK',#             num_blocks=(4, ),#             num_channels=(64, )),#         stage2=dict(#             num_modules=1,#             num_branches=2,#             block='BASIC',#             num_blocks=(4, 4),#             num_channels=(32, 64)),#         stage3=dict(#             num_modules=4,#             num_branches=3,#             block='BASIC',#             num_blocks=(4, 4, 4),#             num_channels=(32, 64, 128)),#         stage4=dict(#             num_modules=3,#             num_branches=4,#             block='BASIC',#             num_blocks=(4, 4, 4, 4),#             num_channels=(32, 64, 128, 256))),#     init_cfg=dict(#         type='Pretrained',#         checkpoint='https://download.openmmlab.com/mmpose/'#         'pretrain_models/hrnet_w32-36af842e.pth'),# ),backbone=dict(type='SCTNet',init_cfg=dict(type='Pretrained',checkpoint=checkpoint_backbone),base_channels=64,spp_channels=128),head=dict(type='HeatmapHead',#in_channels=32,in_channels=256,out_channels=17,deconv_out_channels=None,loss=dict(type='KeypointMSELoss', use_target_weight=True),decoder=codec),test_cfg=dict(flip_test=True,flip_mode='heatmap',shift_heatmap=True,))

4.3.2.4 修改新模型

上述我们分别在复现mmpose框架和我们想要采用的分割开源项目后，可以发现即使统一输入图像为640x640，两个项目的backbone的输出都是不一样的。这里因为基于mmpose框架来改，所以我们肯定是将分割开源项目backbone的输出尽量调整和我们上述复现mmpose中hrnet的输出一致，这样在框架后续中求loss等也就不会报错了。
这是分割开源项目的backbone的输出：

这是mmpose中hrnet的输出：

两个backbone输出的通道数不一致，这会导致Head报错，这里我们对配置文件的Head接受输入的通道数进行修改，将32更改为我们的新backbone的输出通道数256。

所以，现在两个backbone的输出还差宽高不一致，我们这个时候来进一步看一下，Head部分的代码，Head采用’HeatmapHead’，我们定位到mmpose/models/heads/heatmap_heads/heatmap_head.py，发现前向传播forward中代码不多，所以我们选择在Head部分调整backbone的输出。

因为我们的backbone中输出的是list，这里我们要对原始的forward进行修改，我们选用backbone最后层的输出，也就是list[0]，同时我们采用双线性插值调整输出的宽高。

    def forward(self, feats: Tuple[Tensor]) -> Tensor:"""Forward the network. The input is multi scale feature maps and theoutput is the heatmap.Args:feats (Tuple[Tensor]): Multi scale feature maps.Returns:Tensor: output heatmap."""# x = feats[-1]x = feats[0]x = self.deconv_layers(x)x = self.conv_layers(x)x = self.final_layer(x)x = F.interpolate(input=x, scale_factor=2, mode='bilinear')return x

4.4 添加auxiliary_head

上述git中使用了Transformer Block作为辅助head，这里我们也尝试在我们的mmpose框架中进行改进。

https://github.com/xzz777/SCTNet

我们参照这个git，添加auxiliary_head的配置文件。

# model settings
model = dict(type='TopdownPoseEstimator',data_preprocessor=dict(type='PoseDataPreprocessor',mean=[123.675, 116.28, 103.53],std=[58.395, 57.12, 57.375],bgr_to_rgb=True),backbone=dict(type='SCTNet',init_cfg=dict(type='Pretrained',checkpoint=checkpoint_backbone),base_channels=64,spp_channels=128),head=dict(type='HeatmapHead_Revise',# in_channels=32,in_channels=256,out_channels=17,deconv_out_channels=None,loss=dict(type='KeypointMSELoss', use_target_weight=True),decoder=codec),auxiliary_head=[dict(type='VitGuidanceHead',init_cfg=dict(type='Pretrained',checkpoint= checkpoint_teacher),in_channels=256,channels=256,base_channels=64,in_index=2,num_classes=171,loss_decode=dict(type='AlignmentLoss', loss_weight=[3, 15, 15, 15]))],test_cfg=dict(flip_test=True,flip_mode='heatmap',shift_heatmap=True,))

但是，发生了报错：

TypeError: init() got an unexpected keyword argument ‘auxiliary_head’

我们在进入底层环境（这里就是我们对应配置的Anaconda环境）后，我的是：D:\Anaconda3\envs\SCTNet\Lib\site-packages\mmengine\registry\build_functions.py，定位到报错的位置，发现是在：mmpose/models/pose_estimators/topdown.py中

在TopdownPoseEstimator类中，添加缺少了‘auxiliary_head’字典的键。

随后，我们得在mmpose框架中注册一个新的head，你需要在目录 mmpose/models/heads/heatmap_heads下新建一个文件进行定义。

新建的head网络需要继承BaseHead类，其他方面与你继承 nn.Module 来创建没有任何不同。

在完成骨干网络的实现后，你需要使用 MODELS来对其进行注册：

from mmpose.registry import MODELS
from ..base_head import BaseHead
@MODELS.register_module()
class YourNewHead(BaseHead): 

最后，请记得在heatmap_heads/__init__.py 中导入你的新nead网络。