HRNet閱讀筆記及代碼理解

chenatu 發(fā)布于2019-07-31 10:25 / 2930人閱讀

摘要：以高分辨率子網(wǎng)開始作為第一階段，逐個添加高到低分辨率子網(wǎng)以形成更多階段，并且并行連接多分辨率子網(wǎng)。優(yōu)點并行連接高低分辨率子網(wǎng)，而不是像大多數(shù)現(xiàn)有解決方案那樣串聯(lián)連接。我們認為原因是從低分辨率子網(wǎng)上的早期階段提取的低級功能不太有用。

摘要：

大多數(shù)現(xiàn)有方法從由高到低分辨率網(wǎng)絡產(chǎn)生的低分辨率表示中恢復高分辨率表示。相反，本文在整個過程中保持高分辨率的表示。我們將高分辨率子網(wǎng)開始作為第一階段，逐步添加高到低分辨率子網(wǎng)以形成更多階段，并行連接多個子網(wǎng)，每個子網(wǎng)具有不同的分辨率。我們進行重復的多尺度融合，使得高到低分辨率表示可以重復從其他分辨率的表示獲取信息，從而導致豐富的高分辨率表示。因此，預測的關(guān)鍵點熱圖可能更準確，空間更精確。

1. 簡介 1.1 現(xiàn)有方法

(a) 對稱結(jié)構(gòu)，先下采樣，再上采樣，同時使用跳層連接恢復下采樣丟失的信息；
(b) 級聯(lián)金字塔；
(c) 先下采樣，轉(zhuǎn)置卷積上采樣，不使用跳層連接進行數(shù)據(jù)融合；
(d) 擴張卷積，減少下采樣次數(shù)，不使用跳層連接進行數(shù)據(jù)融合；

1.2 HRNet

簡要描述：

HighResolution Net（HRNet），它能夠在整個過程中保持高分辨率表示。以高分辨率子網(wǎng)開始作為第一階段，逐個添加高到低分辨率子網(wǎng)以形成更多階段，并且并行連接多分辨率子網(wǎng)。在整個過程中反復交換并行多分辨率子網(wǎng)絡中的信息來進行重復的多尺度融合。

優(yōu)點：

（a）并行連接高低分辨率子網(wǎng)，而不是像大多數(shù)現(xiàn)有解決方案那樣串聯(lián)連接。因此，我們的方法能夠保持高分辨率而不是通過從低到高的過程恢復分辨率，因此預測的熱圖可能在空間上更精確

（b）大多數(shù)現(xiàn)有的融合方案匯總了低級別和高級別的表示。相反，我們在相同深度和相似水平的低分辨率表示的幫助下執(zhí)行重復的多尺度融合以提升高分辨率表示，反之亦然，導致高分辨率表示對于姿勢估計也是豐富的。因此，我們預測的熱圖可能更準確。個人感覺增加多尺度信息之間的融合是正確的，例如原圖像和模糊圖像進行聯(lián)合雙邊濾波可以得到介于兩者之間的模糊程度的圖像，而RGF濾波就是重復將聯(lián)合雙邊濾波的結(jié)果作為那張模糊的引導圖，這樣得到的結(jié)果會越來越趨近于原圖。此處同樣的道理，不同分辨率的圖像采樣到相同的尺度反復的融合，加之網(wǎng)絡的學習能力，會使得多次融合后的結(jié)果更加趨近于正確的表示。

2. 方法描述 2.1 并行高分辨率子網(wǎng)

2.2 重復多尺度融合

3. 實驗部分 3.1 消融研究 3.1.1 重復多尺度融合

(a) W / o中間交換單元（1個融合）：除最后一個交換單元外，多分辨率子網(wǎng)之間沒有交換；
(b) 僅W /跨階段交換單元（3個融合）：每個階段內(nèi)并行子網(wǎng)之間沒有交換；
(c) W /跨階段和階段內(nèi)交換單元（共8個融合）：這是我們提出的方法；

3.1.2 分辨率保持

所有四個高到低分辨率子網(wǎng)都在開頭添加，深度相同，融合方案與我們的相同。該變體實現(xiàn)了72.5的AP，低于我們的小型網(wǎng)HRNet-W32的73.4 AP。我們認為原因是從低分辨率子網(wǎng)上的早期階段提取的低級功能不太有用。此外，沒有低分辨率并行子網(wǎng)的類似參數(shù)和計算復雜度的簡單高分辨率網(wǎng)絡表現(xiàn)出低得多的性能。

3.1.3 分辨率表示質(zhì)量

檢查從每個分辨率的特征圖估計的熱圖的質(zhì)量。

4. 代碼學習（源碼地址） 4.1 ResNet模塊

雖然很熟悉了，但是還是介紹一下resnet網(wǎng)絡的基本模塊。如下的左圖對應于resnet-18/34使用的基本塊，右圖是50/101/152所使用的，由于他們都比較深，所以有圖相比于左圖使用了1x1卷積來降維。

(a) conv3x3: 沒啥好解釋的，將原有的pytorch函數(shù)固定卷積和尺寸為3重新封裝了一次；

(b) BasicBlock: 搭建上圖左邊的模塊。

(1) 每個卷積塊后面連接BN層進行歸一化；

(2) 殘差連接前的3x3卷積之后只接入BN，不使用ReLU，避免加和之后的特征皆為正，保持特征的多樣；

(3) 跳層連接：兩種情況，當模塊輸入和殘差支路（3x3->3x3）的通道數(shù)一致時，直接相加；當兩者通道不一致時（一般發(fā)生在分辨率降低之后，同分辨率一般通道數(shù)一致），需要對模塊輸入特征使用1x1卷積進行升/降維（步長為2，上面說了分辨率會降低），之后同樣接BN，不用ReLU。

(1) 使用1x1卷積先降維，再使用3x3卷積進行特征提取，最后再使用1x1卷積把維度升回去；

(2) 每個卷積塊后面連接BN層進行歸一化；

(2) 殘差連接前的1x1卷積之后只接入BN，不使用ReLU，避免加和之后的特征皆為正，保持特征的多樣性。

(3) 跳層連接：兩種情況，當模塊輸入和殘差支路（1x1->3x3->1x1）的通道數(shù)一致時，直接相加；當兩者通道不一致時（一般發(fā)生在分辨率降低之后，同分辨率一般通道數(shù)一致），需要對模塊輸入特征使用1x1卷積進行升/降維（步長為2，上面說了分辨率會降低），之后同樣接BN，不用ReLU。

def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
                               bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion,
                                  momentum=BN_MOMENTUM)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

4.2 HighResolutionModule (高分辨率模塊)

當僅包含一個分支時，生成該分支，沒有融合模塊，直接返回；當包含不僅一個分支時，先將對應分支的輸入特征輸入到對應分支，得到對應分支的輸出特征；緊接著執(zhí)行融合模塊。

(a) _check_branches: 判斷num_branches (int) 和 num_blocks, num_inchannels, num_channels (list) 三者的長度是否一致，否則報錯；

(b) _make_one_branch: 搭建一個分支，單個分支內(nèi)部分辨率相等，一個分支由num_blocks[branch_index]個block組成，block可以是兩種ResNet模塊中的一種；
(1) 首先判斷是否降維或者輸入輸出的通道(num_inchannels[branch_index]和 num_channels[branch_index] * block.expansion(通道擴張率))是否一致，不一致使用1z1卷積進行維度升/降，后接BN，不使用ReLU；
(2) 順序搭建num_blocks[branch_index]個block，第一個block需要考慮是否降維的情況，所以多帶帶拿出來，后面1 到 num_blocks[branch_index]個block完全一致，使用循環(huán)搭建就行。此時注意在執(zhí)行完第一個block后將num_inchannels[branch_index重新賦值為  num_channels[branch_index] * block.expansion。


(c)  _make_branches: 循環(huán)調(diào)用_make_one_branch函數(shù)創(chuàng)建多個分支；

(d) _make_fuse_layers:
(1)  如果分支數(shù)等于1，返回None，說明此事不需要使用融合模塊；

(2)  雙層循環(huán)：for i in range(num_branches if self.multi_scale_output else 1):的作用是，如果需要產(chǎn)生多分辨率的結(jié)果，就雙層循環(huán)num_branches 次，如果只需要產(chǎn)生最高分辨率的表示，就將i確定為0。
(2.1) 如果j > i，此時的目標是將所有分支上采樣到和i分支相同的分辨率并融合，也就是說j所代表的分支分辨率比i分支低，2**(j-i)表示j分支上采樣這么多倍才能和i分支分辨率相同。先使用1x1卷積將j分支的通道數(shù)變得和i分支一致，進而跟著BN，然后依據(jù)上采樣因子將j分支分辨率上采樣到和i分支分辨率相同，此處使用最近鄰插值；


(2.2) 如果j = i，也就是說自身與自身之間不需要融合，nothing to do；

(2.3) 如果j < i，轉(zhuǎn)換角色，此時最終目標是將所有分支采樣到和i分支相同的分辨率并融合，注意，此時j所代表的分支分辨率比i分支高，正好和(2.1)相反。此時再次內(nèi)嵌了一個循環(huán)，這層循環(huán)的作用是當i-j > 1時，也就是說兩個分支的分辨率差了不止二倍，此時還是兩倍兩倍往上采樣，例如i-j = 2時，j分支的分辨率比i分支大4倍，就需要上采樣兩次，循環(huán)次數(shù)就是2；
(2.3.1) 當k == i - j - 1時，舉個例子，i = 2,j = 1, 此時僅循環(huán)一次，并采用當前模塊，此時直接將j分支使用3x3的步長為2的卷積下采樣(不使用bias)，后接BN，不使用ReLU；


(2.3.2) 當k != i - j - 1時，舉個例子，i = 3,j = 1, 此時循環(huán)兩次，先采用當前模塊，將j分支使用3x3的步長為2的卷積下采樣(不使用bias)兩倍，后接BN和ReLU，緊跟著再使用(2.3.1)中的模塊，這是為了保證最后一次二倍下采樣的卷積操作不使用ReLU，猜測也是為了保證融合后特征的多樣性；









(e) forward: 前向傳播函數(shù)，利用以上函數(shù)的功能搭建一個HighResolutionModule；
(1) 當僅包含一個分支時，生成該分支，沒有融合模塊，直接返回；

(2) 當包含不僅一個分支時，先將對應分支的輸入特征輸入到對應分支，得到對應分支的輸出特征；緊接著執(zhí)行融合模塊；
(2.1)  循環(huán)將對應分支的輸入特征輸入到對應分支模型中，得到對應分支的輸出特征；
(2.2) 融合模塊：對著這張圖看，很容易看懂。每次多尺度之間的加法運算都是從最上面的尺度開始往下加，所以y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])；加到他自己的時候，不需要經(jīng)過融合函數(shù)的處理，直接加，所以if i == j: y = y + x[j]；遇到不是最上面的尺度那個特征圖或者它本身相同分辨率的那個特征圖時，需要經(jīng)過融合函數(shù)處理再加，所以y = y + self.fuse_layers[i][j](x[j])。最后將ReLU激活后的融合(加法)特征append到x_fuse，x_fuse的長度等于1（單尺度輸出）或者num_branches（多尺度輸出）。

class HighResolutionModule(nn.Module):
    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
                 num_channels, fuse_method, multi_scale_output=True):
        super(HighResolutionModule, self).__init__()
        self._check_branches(
            num_branches, blocks, num_blocks, num_inchannels, num_channels)

        self.num_inchannels = num_inchannels
        self.fuse_method = fuse_method
        self.num_branches = num_branches

        self.multi_scale_output = multi_scale_output

        self.branches = self._make_branches(
            num_branches, blocks, num_blocks, num_channels)
        self.fuse_layers = self._make_fuse_layers()
        self.relu = nn.ReLU(True)

    def _check_branches(self, num_branches, blocks, num_blocks,
                        num_inchannels, num_channels):
        if num_branches != len(num_blocks):
            error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format(
                num_branches, len(num_blocks))
            logger.error(error_msg)
            raise ValueError(error_msg)

        if num_branches != len(num_channels):
            error_msg = "NUM_BRANCHES({}) <> NUM_CHANNELS({})".format(
                num_branches, len(num_channels))
            logger.error(error_msg)
            raise ValueError(error_msg)

        if num_branches != len(num_inchannels):
            error_msg = "NUM_BRANCHES({}) <> NUM_INCHANNELS({})".format(
                num_branches, len(num_inchannels))
            logger.error(error_msg)
            raise ValueError(error_msg)

    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
                         stride=1):
        # ---------------------------(1) begin---------------------------- #
        downsample = None
        if stride != 1 or 
           self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.num_inchannels[branch_index],
                    num_channels[branch_index] * block.expansion,
                    kernel_size=1, stride=stride, bias=False
                ),
                nn.BatchNorm2d(
                    num_channels[branch_index] * block.expansion,
                    momentum=BN_MOMENTUM
                ),
            )
        # ---------------------------(1) end---------------------------- #
        
        # ---------------------------(2) begin---------------------------- #
        layers = []
        layers.append(
            block(
                self.num_inchannels[branch_index],
                num_channels[branch_index],
                stride,
                downsample
            )
        )
        # ---------------------------(2) middle---------------------------- #
        self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
        for i in range(1, num_blocks[branch_index]):
            layers.append(
                block(
                    self.num_inchannels[branch_index],
                    num_channels[branch_index]
                )
            )
        # ---------------------------(2) end---------------------------- #
        return nn.Sequential(*layers)

    def _make_branches(self, num_branches, block, num_blocks, num_channels):
        branches = []

        for i in range(num_branches):
            branches.append(
                self._make_one_branch(i, block, num_blocks, num_channels)
            )

        return nn.ModuleList(branches)

    def _make_fuse_layers(self):
        # ---------------------------(1) begin---------------------------- #
        if self.num_branches == 1:
            return None
        # ---------------------------(1) end---------------------------- #

        
        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        # ---------------------------(2) begin---------------------------- #
        fuse_layers = []
        for i in range(num_branches if self.multi_scale_output else 1):
            fuse_layer = []
            for j in range(num_branches):
                # ---------------------------(2.1) begin---------------------------- #
                if j > i:
                    fuse_layer.append(
                        nn.Sequential(
                            nn.Conv2d(
                                num_inchannels[j],
                                num_inchannels[i],
                                1, 1, 0, bias=False
                            ),
                            nn.BatchNorm2d(num_inchannels[i]),
                            nn.Upsample(scale_factor=2**(j-i), mode="nearest")
                        )
                    )
                # ---------------------------(2.1) end---------------------------- #
                
                # ---------------------------(2.2) begin---------------------------- #
                elif j == i:
                    fuse_layer.append(None)
                # ---------------------------(2.2) end---------------------------- #
                
                # ---------------------------(2.3) begin---------------------------- #
                else:
                    conv3x3s = []
                    for k in range(i-j):
                        # ---------------------------(2.3.1) begin---------------------------- #
                        if k == i - j - 1:
                            num_outchannels_conv3x3 = num_inchannels[i]
                            conv3x3s.append(
                                nn.Sequential(
                                    nn.Conv2d(
                                        num_inchannels[j],
                                        num_outchannels_conv3x3,
                                        3, 2, 1, bias=False
                                    ),
                                    nn.BatchNorm2d(num_outchannels_conv3x3)
                                )
                            )
                        # ---------------------------(2.3.1) end---------------------------- #
                        
                        # ---------------------------(2.3.1) begin---------------------------- #
                        else:
                            num_outchannels_conv3x3 = num_inchannels[j]
                            conv3x3s.append(
                                nn.Sequential(
                                    nn.Conv2d(
                                        num_inchannels[j],
                                        num_outchannels_conv3x3,
                                        3, 2, 1, bias=False
                                    ),
                                    nn.BatchNorm2d(num_outchannels_conv3x3),
                                    nn.ReLU(True)
                                )
                            )
                        # ---------------------------(2.3.1) end---------------------------- #
                    # ---------------------------(2.3) end---------------------------- #
                    fuse_layer.append(nn.Sequential(*conv3x3s))
            fuse_layers.append(nn.ModuleList(fuse_layer))
            # ---------------------------(2) end---------------------------- #

        return nn.ModuleList(fuse_layers)

    def get_num_inchannels(self):
        return self.num_inchannels

    def forward(self, x):
        # ---------------------------(1) begin---------------------------- #
        if self.num_branches == 1:
            return [self.branches[0](x[0])]
        # ---------------------------(1) end---------------------------- #

        # ---------------------------(2) begin---------------------------- #
        # ---------------------------(2.1) begin---------------------------- #
        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])
        # ---------------------------(2.1) end---------------------------- #

        # ---------------------------(2.2) begin---------------------------- #
        x_fuse = []

        for i in range(len(self.fuse_layers)):
            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
            for j in range(1, self.num_branches):
                if i == j:
                    y = y + x[j]
                else:
                    y = y + self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))
        # ---------------------------(2.2) end---------------------------- #
        # ---------------------------(2) end---------------------------- #

        return x_fuse