姿態估計1-07：HR-Net(人體姿態估算)-源碼無死角解析（3）-模型總體結構

以下鏈接是個人關於HR-Net(人體姿態估算) 所有見解，如有錯誤歡迎大家指出，我會第一時間糾正。有興趣的朋友可以加微信：a944284742相互討論技術。若是幫助到了你什麼，一定要記得點贊！因爲這是對我最大的鼓勵。
姿態估計1-00：HR-Net(人體姿態估算)-目錄-史上最新無死角講解

前言

根據前面的博客，我們已經知道了數據讀取，以及數據預處理的過程，總的來說，就是讀取coco數據的標籤信息，然後轉化爲heatmap。接下來我們去分析構建模型的總體思路。那麼我們下面就開始吧，我們可以tools/train.py中找到如下代碼：

    # 根據配置文件構建網絡
    print('models.'+cfg.MODEL.NAME+'.get_pose_net')
    model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
        cfg, is_train=True
    )

該處，就是構建網絡的代碼，其最後中會調用/lib/models/pose_hrnet.py的def get_pose_net(cfg, is_train, **kwargs):函數：

def get_pose_net(cfg, is_train, **kwargs):
    model = PoseHighResolutionNet(cfg, **kwargs)

    if is_train and cfg['MODEL']['INIT_WEIGHTS']:
        model.init_weights(cfg['MODEL']['PRETRAINED'])

    return model

接下來我們就是要重點分析其中的PoseHighResolutionNet。

PoseHighResolutionNet

代碼註釋如下，這裏只是註釋了比較重要的一部分(大家大致瀏覽以下即可，後面有代碼領讀)：

class PoseHighResolutionNet(nn.Module):

    def __init__(self, cfg, **kwargs):
        self.inplanes = 64
        extra = cfg['MODEL']['EXTRA']
        super(PoseHighResolutionNet, self).__init__()

        # stem net,進行一系列的卷積操作,獲得最初始的特徵圖N11
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self._make_layer(Bottleneck, 64, 4)



        # 獲取stage2的相關配置信息
        self.stage2_cfg = extra['STAGE2']
        # num_channels=[32,64],num_channels表示輸出通道,最後的64是新建平行分支N2的輸出通道數
        num_channels = self.stage2_cfg['NUM_CHANNELS']
        # 這裏的block爲Bottleneck,在論文中有提到,第一個stage到第二個stage變換時,使用Bottleneck
        block = blocks_dict[self.stage2_cfg['BLOCK']]
        # block.expansion默認爲1,num_channels表示輸出通道[32,64]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))
        ]
        # 這裏會生成新的平行分N2支網絡,即N11-->N21,N22這個過程
        # 同時會對輸入的特徵圖x進行通道變換(如果輸入輸出通道書不一致)
        self.transition1 = self._make_transition_layer([256], num_channels)
        # 對平行子網絡進行加工,讓其輸出的y,可以當作下一個stage的輸入x,
        # 這裏的pre_stage_channels爲當前stage的輸出通道數,也就是下一個stage的輸入通道數
        # 同時平行子網絡信息交換模塊,也包含再其中
        self.stage2, pre_stage_channels = self._make_stage(
            self.stage2_cfg, num_channels
        )



        # 獲取stage3的相關配置信息
        self.stage3_cfg = extra['STAGE3']
        # num_channels=[32,64,128],num_channels表示輸出通道,最後的128是新建平行分支N3的輸出通道數
        num_channels = self.stage3_cfg['NUM_CHANNELS']
        # 這裏的block爲BasicBlock,在論文中有提到,除了第一個stage到第二個stage變換時使用Bottleneck,其餘的都是使用BasicBlock
        block = blocks_dict[self.stage3_cfg['BLOCK']]
        # block.expansion默認爲1,num_channels表示輸出通道[32,64,128]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))
        ]
        # 這裏會生成新的平行分支N3網絡,即N22-->N32,N33這個過程
        # 同時會對輸入的特徵圖x進行通道變換(如果輸入輸出通道書不一致)
        self.transition2 = self._make_transition_layer(
            pre_stage_channels, num_channels)
        # 對平行子網絡進行加工,讓其輸出的y,可以當作下一個stage的輸入x,
        # 這裏的pre_stage_channels爲當前stage的輸出通道數,也就是下一個stage的輸入通道數
        # 同時平行子網絡信息交換模塊,也包含再其中
        self.stage3, pre_stage_channels = self._make_stage(
            self.stage3_cfg, num_channels)



        # 獲取stage4的相關配置信息
        self.stage4_cfg = extra['STAGE4']
        # num_channels=[32,64,128,256],num_channels表示輸出通道,最後的256是新建平行分支N4的輸出通道數
        num_channels = self.stage4_cfg['NUM_CHANNELS']
        # 這裏的block爲BasicBlock,在論文中有提到,除了第一個stage到第二個stage變換時使用Bottleneck,其餘的都是使用BasicBlock
        block = blocks_dict[self.stage4_cfg['BLOCK']]
        # block.expansion默認爲1,num_channels表示輸出通道[32,64,128]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))
        ]
        # 這裏會生成新的平行分支N4網絡,即N33-->N43,N44這個過程
        # 同時會對輸入的特徵圖x進行通道變換(如果輸入輸出通道書不一致)
        self.transition3 = self._make_transition_layer(
            pre_stage_channels, num_channels)
        # 對平行子網絡進行加工,讓其輸出的y,可以當作下一個stage的輸入x,
        # 這裏的pre_stage_channels爲當前stage的輸出通道數,也就是下一個stage的輸入通道數
        # 同時平行子網絡信息交換模塊,也包含再其中
        self.stage4, pre_stage_channels = self._make_stage(
            self.stage4_cfg, num_channels, multi_scale_output=False)


        # 對最終的特徵圖混合之後進行一次卷積, 預測人體關鍵點的heatmap
        self.final_layer = nn.Conv2d(
            in_channels=pre_stage_channels[0],
            out_channels=cfg['MODEL']['NUM_JOINTS'],
            kernel_size=extra['FINAL_CONV_KERNEL'],
            stride=1,
            padding=1 if extra['FINAL_CONV_KERNEL'] == 3 else 0
        )

        # 預測人體關鍵點的heatmap
        self.pretrained_layers = extra['PRETRAINED_LAYERS']


    def forward(self, x):

        # 經過一系列的卷積, 獲得初步特徵圖,總體過程爲x[b,3,256,192]-->x[b,256,64,48]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.layer1(x)




        # 對應論文中的stage2
        # 其中包含了創建分支的過程,即 N11-->N21,N22 這個過程
        # N22的分辨率爲N21的二分之一,總體過程爲:
        # x[b,256,64,48] ---> [1, 32, 64, 48]  因爲通道數不一致,通過卷積進行通道數變換
        #                     [1, 64, 32, 24]  通過新建平行分支生成
        x_list = []
        for i in range(self.stage2_cfg['NUM_BRANCHES']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)

        # 總體過程如下(經過一些卷積操作,但是特徵圖的分辨率和通道數都沒有改變)：
        # x[1, 32, 64, 48] --->  y[1, 32, 64, 48]
        # x[1, 64, 32, 24] --->  y[1, 64, 32, 24]
        y_list = self.stage2(x_list)





        # 對應論文中的stage3
        # 其中包含了創建分支的過程,即 N22-->N32,N33 這個過程
        # N33的分辨率爲N32的二分之一,
        # y[1, 32, 64, 48] ---> x[1, 32,  64, 48]   因爲通道數一致,沒有做任何操作
        # y[1, 64, 32, 24] ---> x[1, 64,  32, 24]   因爲通道數一致,沒有做任何操作
        #                       x[1, 128, 16, 12]   通過新建平行分支生成
        x_list = []
        for i in range(self.stage3_cfg['NUM_BRANCHES']):
            if self.transition2[i] is not None:
                x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])


        # 總體過程如下(經過一些卷積操作,但是特徵圖的分辨率和通道數都沒有改變)：
        # x[1, 32, 64, 48] ---> x[1, 32,  64, 48]
        # x[1, 32, 32, 24] ---> x[1, 64,  32, 24]
        # x[1, 64, 16, 12] ---> x[1, 128, 16, 12]
        y_list = self.stage3(x_list)





        # 對應論文中的stage4
        # 其中包含了創建分支的過程,即 N33-->N43,N44 這個過程
        # N44的分辨率爲N43的二分之一
        # y[1, 32,  64, 48] ---> x[1, 32,  64, 48]  因爲通道數一致,沒有做任何操作
        # y[1, 64,  32, 24] ---> x[1, 64,  32, 24]  因爲通道數一致,沒有做任何操作
        # y[1, 128, 16, 12] ---> x[1, 128, 16, 12]  因爲通道數一致,沒有做任何操作
        #                        x[1, 256, 8,  6 ]  通過新建平行分支生成

        x_list = []
        for i in range(self.stage4_cfg['NUM_BRANCHES']):
            if self.transition3[i] is not None:
                x_list.append(self.transition3[i](y_list[-1]))
            else:
                x_list.append(y_list[i])

        # 進行多尺度特徵融合
        # x[1, 32,  64, 48] --->
        # x[1, 64,  32, 24] --->
        # x[1, 128, 16, 12] --->
        # x[1, 256, 8,  6 ] --->   y[1, 32,  64, 48]
        y_list = self.stage4(x_list)

        # y[1, 32, 64, 48] --> x[1, 17, 64, 48]
        x = self.final_layer(y_list[0])

        return x

其上爲兩個部分，分別爲初始化過程,以及前向傳播的過程。

forward

對於前線傳播過程，大家看了註釋之後應該是很清楚很明白了,主要對應論文中的如下過程：
$N_{11} \rightarrow N_{21}\rightarrow N_{31} \rightarrow N_{41} \\ \quad \quad \searrow N_{22} \rightarrow N_{32}\rightarrow N_{42} \\ \quad \quad \quad \quad \quad \searrow N_{33} \rightarrow N_{43} \\ \quad \quad \quad \quad \quad \quad \quad \quad \searrow N_{44}$

通過代碼我們可以很明顯的看到，最終我們獲得一個[1, 17, 64, 48]大小的heatmap，這就是我們最終想要的結果，其上的每個$N_{XX}$可以分成兩個重要的模塊，分別爲 self.transition_x以及 self.transition_x，其再初始化函數中構建。

init 函數

def __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 主要是創建新的平行子分支網絡, _make_stage是爲了構建論文中平行子網絡信息交流的模塊。具體的實現過程,在下篇博客中爲大家講解。