Stacked Hourglass筆記&源碼（一）網絡結構

前言：

看了很多關於Stacked Hourglass網絡模型的解析的博文，大多數都是隻是對模型結構的分析，或者是對論文的翻譯，但是自己根據論文以及去復現的時候卻遇到很多問題，比如計算loss時要如何計算，acc如何計算，生成的heatmap如何轉爲關鍵點等等小問題。這些問題十分重要，但是在論文以及博文中並沒有詳細的講到如何處理。因此自己根據在復現的經歷，寫這一系列博客進行記錄，希望能幫到需要的朋友。

• 該系列圍繞下圖進行記錄，初步打算分 網絡結構、訓練細節、運行demo這三個主題進行編寫，後期可能會隨時改動。
  

• 當前已完成至訓練細節的loss計算，有些細節還未扣清楚，之後會一點點更新

• 我的復現進度：mxnet版Stacked Hourglass

概覽

模形都由子模塊一點點拼湊成一個大網絡，該篇也是打算從子模塊講起，再講講子模塊組成的大一點點的模塊，再講大一點點的模塊構成的整體網絡。整體模塊的組成結構如下圖

當然這只是大概的結構，在各個大小模塊的組成之間還有一些小的細節。下面開始各個模塊的介紹。

Redidual

我在這裏貼了倆張圖，第二張是我自己畫的，感覺會比上一張詳細，但是不如上一張明瞭。

1. 圖一中藍色的代表BatchNormal 紫色的代表Relu
2. 圖二中的參數是通過官方源碼獲取的，所以跟着圖片設置就好啦

總體上來看，Redidual模塊分爲倆條支路，一條是convBlock，一條是skipLayer，很顯然就是殘差網絡的那個思想。convBlock對輸入的數據進行深度的卷積，以便獲取深層次的圖像信息，最後再將原來輸入的數據經過skipLayer進行element add並輸出，這樣可以將深層次的信息和原始信息進行相結合，獲得更好的效果。

對於convBlock，無腦照着參數卷積就是了
對於skipLayer，當輸入數據的channel與Redidual模塊要求的output_channel不一致的時候要使用1*1的conv卷積核進行卷積，以便後面的element_add能加的起來

以下是mxnet版本的代碼：

class Residual(nn.HybridBlock):

    def __init__(self, in_channels, out_channels, **kwargs):
        super(Residual, self).__init__(**kwargs)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.residual_conv = nn.HybridSequential()
        self.residual_skip = nn.HybridSequential()

        with self.residual_conv.name_scope():
            # 卷積路
            self.residual_conv.add(nn.BatchNorm())
            self.residual_conv.add(nn.Activation('relu'))
            self.residual_conv.add(nn.Conv2D(self.out_channels // 2, (1, 1)))
            self.residual_conv.add(nn.BatchNorm())
            self.residual_conv.add(nn.Activation('relu'))
            self.residual_conv.add(nn.Conv2D(self.out_channels // 2, (3, 3), (1, 1), (1, 1)))
            self.residual_conv.add(nn.BatchNorm())
            self.residual_conv.add(nn.Activation('relu'))
            self.residual_conv.add(nn.Conv2D(self.out_channels, (1, 1)))

        # 連接路
        if not self.in_channels == self.out_channels:
            with self.residual_skip.name_scope():
                self.residual_skip.add(nn.Conv2D(self.out_channels, (1, 1)))
        
    def hybrid_forward(self, F, x):    
        temp_x = x
        x = self.residual_conv(x)
        if not self.in_channels == self.out_channels:
            x = x + self.residual_skip(temp_x)
        else:
            x = x + temp_x
        return x

Hourglass

這個模塊就比較靈活了，是整個模型的核心，是由Redidual爲基礎模塊構成。這個模塊還有一個階數n，不同的階數有不同的特徵。

一階Hourglass長這樣

二階Hourglass長這樣

四階長這樣

各階區別：

通過觀察不難發現，一階二階以及多階的區別就在於虛線框內裝的不一樣
我在閱讀官方源碼的時候注意到以下幾點：

1. 綠色的代表Redidual模塊，箭頭朝下是MaxPool，箭頭朝上是UpSamplingNearest
2. 上面幾張圖的大部分Redidual塊（上方的SkipLayer、下采樣部分前後）都是3*Redidual構成，但是在官方代碼中，這個的數量通過一個叫opt.nModules參數進行設置的，圖片中爲3，官方源碼爲1，本文以及復現代碼按照官方源碼的1來設置（下面的圖會略微不同）。
3. 官方源碼中階數由參數n控制，默認爲4階。由遞歸方法去構造n階的Hourglass（具體看下面源碼）。

下圖是我自己根據官方源碼畫的詳細一點的Hourglass模塊

看圖會發現，這裏也有殘差的思想（跳級結構）(也可看做Skiplayer 與 convBlock)，這個思想很無敵啊。。。
其次就是Hourglass（漏斗）的體現就在於先MaxPool後UpSampling,
通過MaxPool將feature map縮小，再UpSampling進行擴大（編碼-解碼那一套）

用圖片畫出來就粗略的長這樣（下圖由2個Hourglass 組成，不含內部構造）：

4階的詳細的長這樣（單個Hourglass以及其內部構造 4階）：

下面是mxnet源碼

1. 我把Hourglass命名爲HourGlassBlock模塊
2. args.nModules參數設置Redidual的數量，根據官方這裏設置爲1

class HourGlassBlock(nn.HybridBlock):

    def __init__(self, n, in_channels, **kwargs):
        '''
        args:
            n:              當前HourGlass所在的階數
            in_channels:    當前HourGlass輸入的channels
        '''
        super(HourGlassBlock, self).__init__(**kwargs)
        self.n = n
        self.in_channels = in_channels

        with self.name_scope():
            # Upper branch
            self.up1 = nn.HybridSequential()
            for _ in range(args.nModules):
                self.up1.add(Residual(self.in_channels, self.in_channels))

            # Lower branch
            self.low1_MaxPool = nn.MaxPool2D((2, 2), (2, 2))
            self.low1 = nn.HybridSequential()
            for _ in range(args.nModules):
                self.low1.add(Residual(self.in_channels, self.in_channels))

			# 遞歸生成上文圖中虛線框內的子階Hourglass
            if self.n > 1:
                self.low2 = HourGlassBlock(self.n - 1, self.in_channels)
            else:
                self.low2 = nn.HybridSequential()
                for _ in range(args.nModules):
                    self.low2.add(Residual(self.in_channels, self.in_channels))

            self.low3 = nn.HybridSequential()
            for _ in range(args.nModules):
                self.low3.add(Residual(self.in_channels, self.in_channels))

    def hybrid_forward(self, F, x):
        up1 = self.up1(x)	# SkipLayer
		# ConvBlock
        x = self.low1_MaxPool(x)
        x = self.low1(x)
        x = self.low2(x)
        x = self.low3(x)

        up2 = F.UpSampling(x, scale=2, sample_type="nearest")

        return up1 + up2

Linear

在講完整網絡之前要提一下這個模塊，該模塊主要由1*1的Conv與BN以及Relu組成，緊跟在Hourglass之後

class Lin(nn.HybridBlock):

    def __init__(self, numOut, **kwargs):
        super(Lin, self).__init__(**kwargs)
        self.numOut = numOut
        self.lin = nn.HybridSequential()

        with self.lin.name_scope():
            self.lin.add(nn.Conv2D(numOut, 1))
            self.lin.add(nn.BatchNorm())
            self.lin.add(nn.Activation('relu'))

    def hybrid_forward(self, F, x):
        return self.lin(x)

完整網絡

現在到具體的完整的網絡

先看看整個網絡大致長啥樣

一張圖片輸入，經過一些簡單的卷積模塊，再通過N個Hourglass模塊，最終得到輸出

下面看看N個Hourglass模塊之前的卷積模塊

1. 輸入的圖片是HW3 (尺寸默認是256*256)
2. 在圖片經過N個n階的Hourglass模塊之前是有一些處理模塊的，位置：圖中從左到右紅色4個模塊

接下來看看封裝的Hourglass模塊

N個n階的Hourglass模塊並不是簡簡單單的拼在一起，每兩個Hourglass之間還有一些卷積模塊（也就是說單個Hourglass是被一些卷積模塊封裝起來的，如下圖），以及論文中提到的“中繼監督”也在這裏運用。下面詳細說說這裏的細節

1. 這裏的輸入指上一個相同的Hourglass模塊或者上文說的Hourglass模塊之前是有一些處理模塊
2. 在這裏有三條SkipLayer，通過element_add相加在一起
3. 在我的代碼中，用了一個[]存放N個out(heatmap)，中繼監督就是這裏體現，存放每一個Hourglass模塊的輸出，N個out都用作loss計算
4. 關於這個out，首先要說明下N在論文中使用的是8（代碼中對應是nStack），那麼out中便有8個heatmap
5. 因爲訓練的數據集MPII一共有(16個關節點)，所以每個heatmap的shape是(16, x, x, x)，一個heatmap對應一個關節點，類似這樣
   
   6.用N個heatmap去計算loss時候，或者計算acc的時候，很顯然要把標註文件也轉爲對應的16個heatmap，然後進行訓練，這個打算在下一篇中講

具體代碼如下：

class Hourglass(nn.HybridBlock):

    def __init__(self, **kwargs):
        super(Hourglass, self).__init__(**kwargs)
        self.out = []

        # HourglassBlock模塊之前的圖片處理模塊
        self.preprocess = nn.HybridSequential(prefix="pre")
        with self.preprocess.name_scope():
            self.preprocess.add(nn.Conv2D(64, 7, (2, 2), (3, 3)))
            self.preprocess.add(nn.BatchNorm())
            self.preprocess.add(nn.Activation("relu"))
            self.preprocess.add(Residual(64, 128))
            self.preprocess.add(nn.MaxPool2D((2, 2), (2, 2)))
            self.preprocess.add(Residual(128, 128))
            self.preprocess.add(Residual(128, args.nFeats))

        # HourglassBlock模塊
        self.hourglass_blocks = nn.HybridSequential(prefix="hg")
        with self.hourglass_blocks.name_scope():
            for _ in range(args.nStack):
                hourglass_block = nn.HybridSequential()
                hourglass_block.add(HourGlassBlock(4, args.nFeats))	# args.nFeats = 256
                for _ in range(args.nModules):	# args.nModules = 1
                    hourglass_block.add(Residual(args.nFeats, args.nFeats))
                hourglass_block.add(Lin(args.nFeats))
                hourglass_block.add(nn.Conv2D(args.nJoints, (1, 1), (1, 1), (0, 0)))	# args.nJoints = 16 數據集16個關節點 這層之後可以生成out
                self.hourglass_blocks.add(hourglass_block)
            self.conv1 = nn.Conv2D(args.nFeats, (1, 1), (1, 1), (0, 0))
            self.conv2 = nn.Conv2D(args.nFeats, (1, 1), (1, 1), (0, 0))

    def hybrid_forward(self, F, x):
        x = self.preprocess(x)
        for i in range(args.nStack):
            temp_x = x
            x = self.hourglass_blocks[i](x)
            self.out.append(x)

            if i < args.nStack:
                x1 = self.conv1(x)
                x2 = self.conv2(x)
                x = temp_x + x1 + x2

        return self.out

參考資料

• 官方源碼
• pytorch版本
• keras版本
• 論文地址（國內查看比較方便的平臺）

如有錯誤，還請指正