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

复现open-mmlab的mmpose详细细节

1.配置环境
2.数据处理
3.训练
4.改进mmpose
- 4.1 快速调试技巧
- 4.2 快速定位
- 4.3 改进backbone
- - 4.3.1 使用说明
  - 4.3.2 改进案例
  - - 4.3.2.1 复现mmpose原配置文件
    - 4.3.2.2 复现开源项目
    - 4.3.2.3 修改配置文件
    - 4.3.2.4 修改新模型
- 4.4 添加auxiliary_head
- - 4.4.1 添加auxiliary_head配置
  - 4.4.2 注意事项

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 the
        output 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

4.4.1 添加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网络。
在这里插入图片描述

4.4.2 注意事项

可能会报错：

No module named ‘mmcv.cnn.utils.weight_init’

改成：

from mmengine.model import BaseModule, constant_init, kaiming_init, trunc_normal_init, trunc_normal_init, normal_init

由于，继承了BaseHead，
在这里插入图片描述
BaseHead里有3个abstractmethod，抽象方法在基类中没有实现，子类必须实现这些方法，否则子类也会被视为抽象类，无法实例化。

在这里插入图片描述
如果我们暂时没有使用到这3个方法，可以重写出来，用pass带过：