姿态估计1-08：HR-Net(人体姿态估算)-源码无死角解析（4）-平行分支，信息交流模块构建

以下链接是个人关于HR-Net(人体姿态估算) 所有见解，如有错误欢迎大家指出，我会第一时间纠正。有兴趣的朋友可以加微信：a944284742相互讨论技术。若是帮助到了你什么，一定要记得点赞！因为这是对我最大的鼓励。
姿态估计1-00：HR-Net(人体姿态估算)-目录-史上最新无死角讲解

前言

通过上篇博客，详细介绍了ib/models/pose_hrnet.py中类 PoseHighResolutionNet 的如下函数：

    def __init__(self, cfg, **kwargs):
    	......
    def forward(self, x):
    	......

到最后，我们发现__init__中调用了：

    def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer):
    def _make_stage(self, layer_config, num_inchannels,multi_scale_output=True):

相对来说，这两个函数是比较复杂，同时也是比较核心的函数，下面我们对其进行分析。

_make_transition_layer

    def _make_transition_layer(
            self, num_channels_pre_layer, num_channels_cur_layer):
        """
        :param num_channels_pre_layer: 上一个stage平行网络的输出通道数目,为一个list,
            stage=2时, num_channels_pre_layer=[256]
            stage=3时, num_channels_pre_layer=[32,64]
            stage=4时, num_channels_pre_layer=[32,64,128]
        :param num_channels_cur_layer:
            stage=2时, num_channels_cur_layer = [32,64]
            stage=3时, num_channels_cur_layer = [32,64,128]
            stage=4时, num_channels_cur_layer = [32,64,128,256]
        """
        num_branches_cur = len(num_channels_cur_layer)
        num_branches_pre = len(num_channels_pre_layer)

        transition_layers = []
        # 对stage的每个分支进行处理
        for i in range(num_branches_cur):
            # 如果不为最后一个分支
            if i < num_branches_pre:
                # 如果当前层的输入通道和输出通道数不相等,则通过卷积对通道数进行变换
                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
                    transition_layers.append(
                        nn.Sequential(
                            nn.Conv2d(
                                num_channels_pre_layer[i],
                                num_channels_cur_layer[i],
                                3, 1, 1, bias=False
                            ),
                            nn.BatchNorm2d(num_channels_cur_layer[i]),
                            nn.ReLU(inplace=True)
                        )
                    )
                # 如果当前层的输入通道和输出通道数相等,则什么都不做
                    transition_layers.append(None)

            # 如果为最后一个分支,则再新建一个分支（该分支分辨率会减少一半）
            else:
                conv3x3s = []
                for j in range(i+1-num_branches_pre):
                    inchannels = num_channels_pre_layer[-1]
                    outchannels = num_channels_cur_layer[i] \
                        if j == i-num_branches_pre else inchannels
                    conv3x3s.append(
                        nn.Sequential(
                            nn.Conv2d(
                                inchannels, outchannels, 3, 2, 1, bias=False
                            ),
                            nn.BatchNorm2d(outchannels),
                            nn.ReLU(inplace=True)
                        )
                    )
                transition_layers.append(nn.Sequential(*conv3x3s))

        return nn.ModuleList(transition_layers)

_make_stage

    def _make_stage(self, layer_config, num_inchannels,
                    multi_scale_output=True):

        """
        当stage=2时： num_inchannels=[32,64]           multi_scale_output=Ture
        当stage=3时： num_inchannels=[32,64,128]       multi_scale_output=Ture
        当stage=4时： num_inchannels=[32,64,128,256]   multi_scale_output=False
        """

        # 当stage=2,3,4时,num_modules分别为：1,4,3
        # 表示HighResolutionModule（平行之网络交换信息模块）模块的数目
        num_modules = layer_config['NUM_MODULES']

        # 当stage=2,3,4时,num_branches分别为：2,3,4,表示每个stage平行网络的数目
        num_branches = layer_config['NUM_BRANCHES']

        # 当stage=2,3,4时,num_blocks分别为：[4,4], [4,4,4], [4,4,4,4],
        # 表示每个stage blocks(BasicBlock或者BasicBlock)的数目
        num_blocks = layer_config['NUM_BLOCKS']

        # 当stage=2,3,4时,num_channels分别为：[32,64],[32,64,128],[32,64,128,256]
        # 在对应stage, 对应每个平行子网络的输出通道数
        num_channels = layer_config['NUM_CHANNELS']

        # 当stage=2,3,4时,分别为：BasicBlock,BasicBlock,
        block = blocks_dict[layer_config['BLOCK']]

        # 当stage=2,3,4时,都为：SUM,表示特征融合的方式
        fuse_method = layer_config['FUSE_METHOD']


        modules = []
        # 根据num_modules的数目创建HighResolutionModule
        for i in range(num_modules):
            # multi_scale_output is only used last module
            # multi_scale_output 只被用再最后一个HighResolutionModule
            if not multi_scale_output and i == num_modules - 1:
                reset_multi_scale_output = False
            else:
                reset_multi_scale_output = True

            # 根据参数,添加HighResolutionModule到
            modules.append(
                HighResolutionModule(
                    num_branches, # 当前stage平行分支的数目
                    block, # BasicBlock,BasicBlock
                    num_blocks, # BasicBlock或者BasicBlock的数目
                    num_inchannels,# 输入通道数目
                    num_channels, # 输出通道数
                    fuse_method, # 通特征融合的方式
                    reset_multi_scale_output # 是否使用多尺度方式输出
                )
            )
            # 获得最后一个HighResolutionModule的输出通道数
            num_inchannels = modules[-1].get_num_inchannels()

        return nn.Sequential(*modules), num_inchannels

总结

我相信通过上面的注解，大家对两个函数应该算是比较了解了。可以很明显的看到，其中还有一个核心就是HighResolutionModule模块，也就是信息交流模块的核心关键，下篇博客会专门对其进行讲解，欢迎观看。