ps:轉載請注明出處，謝謝。
以下簡稱HRnet
這篇論文我拖更了好久，早在半年前我就說我要更新和這篇文獻相關的代碼研讀，一直是懶，然后代碼太長，分析代碼真的要有決心+耐心+毅力，不然的話很容易放棄的，一件事情你做了百分之99就等同於沒有做，行百里者半九十，就是這個道理，希望所有在這個領域內的小白通過閱讀文獻，編寫代碼來提升自己，相信你自己，你挺棒的。
HRnet由最基本的三種塊構成。第一種是普通的3x3的卷積它的結構如下。
它的代碼如下
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)
第二種是BasicBlock，它的結構如下。
當inchannels和outchannels不想等時就進行將采樣。它的代碼如下
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
第三種結構是三層的殘差塊，結構如下圖。這個結構里面有一個參數叫做expansion的參數，這個參數用來控制卷積的輸入輸出通道數。
其中BN層和RELU層我就不單獨的畫出來了，想看可以去原文中找相應的代碼，這一部分的代碼如下
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
就和正常三層卷積的殘差塊是一樣的道理。接下來就是高分辨率模塊。首先會看到高分辨率模塊的參數列表如下
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(False)
check_branches()這個函數這個函數的作用是檢查，在高分辨率模塊中num_branches(int類型)，和len(num_inchannels(里面的元素是int))，和len(num_channels(里面的元素是int))它們三個的值是否相等，如果不想等就報出異常。那么這三個變量是什么意思呢。我們首先看一下高分辨率模塊的圖。
這個圖里面我畫出的這個部分就是一個完整的高分辨率模塊，num_branches代表的是有幾個分支，就是融合的時候(交叉很多線那個地方)，有幾條線指向一組featuremaps，（不算多出來的新分支，那是新的stage要考慮的問題），那么我這個圖里面是每個featuremaps組有兩條線指向它，因此num_branches=2，num_inchannels和num_channels都是列表，他們表示的是featuremaps的輸入通道數和輸出通道數，因為同一個分支上的featuremaps的通道數是一致的，因此有幾個分支(num_branches)，len（num_inchannels）和len(num_channels)就是幾。同樣就有幾種尺度的featuremaps的融合，這一點后面也會說到。下面是check_branches部分的代碼。
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)
下面一個函數是_make_one_branch，代碼如下
def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1): 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), ) layers = [] layers.append(block(self.num_inchannels[branch_index], num_channels[branch_index], stride, downsample)) 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])) return nn.Sequential(*layers)
它的作用就是創建一個新的分支，就是
我畫出的這個部分，因為HRnet的所有塊都是基於兩種殘差塊的，因此只有每個分支的第一個塊是特殊的，我們首先要判斷的是inchannels和channels*expansion是否相等，如果不相等(針對的是Bottleneck塊)，那么就要進行downsample的操作，將通道調整到一致，因為每個branch里面有num_blocks[branch_index]個塊，除了第一個塊不同，其他的全部相同，因此采用一個for循環，從1到num_blocks[branch_index]逐漸生成基本塊即可。
make_branches函數是看看每個stage里面有多少branch，然后有幾個就調用幾次_make_one_branch函數，這部分代碼就不再贅述了，就一個for循環，簡單。
重點重點重點來了，我把字放大來講，這部分實在實在是不好理解，如果你覺得你看起來沒問題的話，那你不用看我嘮叨了。
上面我講解了各個參數的含義，fuse_layer有個地方特別難懂這也是我第一遍看沒有看懂的原因，那么它說的到底是個啥意思呢，請看下面的圖。
然后我們看一下代碼如下
def _make_fuse_layers(self): if self.num_branches == 1: return None num_branches = self.num_branches num_inchannels = self.num_inchannels fuse_layers = [] for i in range(num_branches if self.multi_scale_output else 1): fuse_layer = [] for j in range(num_branches): 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], momentum=BN_MOMENTUM), nn.Upsample(scale_factor=2**(j-i), mode='nearest'))) elif j == i: fuse_layer.append(None) else: conv3x3s = [] for k in range(i-j): 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, momentum=BN_MOMENTUM))) 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, momentum=BN_MOMENTUM), nn.ReLU(False))) fuse_layer.append(nn.Sequential(*conv3x3s)) fuse_layers.append(nn.ModuleList(fuse_layer)) return nn.ModuleList(fuse_layers)
首先要注意的是fuselayers里面是不包含我畫叉的那一條紅色的線的，只包含粉色線和藍色線的操作，那雙重循環里面的i代表什么呢，i代表的當前融合的branch，上面的圖我畫出了當i=0時，所有的featuremaps都融合到0這個分支的featuremaps上面去，j代表組成融合的featuremaps所對應的branchindex，那么這時候要分三種情況討論。
第一種情況：j>i
此時j所在分支的featuremaps的分辨率比i要小，通道數要多，那此時需要先使用卷積對其進行通道的改變，然后進行上采樣，上采樣因子即scale_factor的大小是2的(j-i)次方，比如j-i=1時那此時j就在i下面一個分支，他們倆的分辨率就差2倍，如果j-i=2，說明j在i下面兩個分支的位置，此時分辨率相差四倍，因為上采樣因子時2的2次方，就是4.
第二種情況，j=i時，j=i時說明要參與融合的分支j和目標分支i在同一個branch上面，因此什么都不用做
第三種情況j<i,那此時j所在分支的分辨率比i所在的目標分支的分辨率要大，因此要進行降采樣，就是改變stride來采樣，那作者引用了一個參數k代表降采樣的次數k，k 的范圍在[0,i-j-1]的這個閉區間內。k每進行循環加1的時候就進行一次卷積改變通道並且降采樣。
當k=i-j-1的時候說明此時 j 分支就在 i 分支的上面一個分支，因此直接輸出通道就是 i 所在分支的輸入通道即可，同時令stride=2改變featuremaps的大小。當k！=i-j-1的時候說明 j 所在分支的分辨率比 i 所在分支的分辨率高出不止2倍，例如 j=0，i=3時，那么令outchannel=in channel[j],這個我也不知道為啥這么做哈哈哈，因為這樣改變通道沒有意義，不如到通道一直不變到最后再變成目標分支 i 的通道呢，我覺得很奇怪。因為當他們的分辨率只差兩倍的時候通道仍相差四倍，然后k=i-j-1直接變到i所在通道數了，很迷惑不知道為什么作者這么做。搞得很復雜。
那么fuselayers最后長啥樣呢，它其實是這樣 fuse_layers[fuse_layer0[j=0 to i=0,j=1 to i=0,j=2 to i=0],fuse_layer1[j=0 to i=1,j=1 to i=1,j=2 to i=1],fuse_layer2[j=0 to i=2,j=1 to i=2,j=2 to i=2]],就是這樣一個列表，是個二維列表(其實就是數組，python沒有數組就用列表代替)，列表中的每一個to操作都是一個sequential或者none，代表j分支到目標i分支的操作。接下來就是forward了，這部分很簡單這里不再贅述，代碼如下。
def forward(self, x): if self.num_branches == 1: return [self.branches[0](x[0])] for i in range(self.num_branches): x[i] = self.branches[i](x[i]) 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)) return x_fuse
剛才那個圖上面畫叉的那條紅線不在fuse_layers中，它在transition_layers中，代碼中有體現，它的作用是每當一個fuse_layers產生時會生成新的分支，新的分支的輸入源於它上一個stage的上一層branch的舊的分支，因此需要額外的考慮一下，不過代碼不是很難理解，看看就懂了。代碼如下
def _make_transition_layer( self, num_channels_pre_layer, num_channels_cur_layer): num_branches_cur = len(num_channels_cur_layer) num_branches_pre = len(num_channels_pre_layer) transition_layers = [] 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], momentum=BN_MOMENTUM), nn.ReLU(inplace=True))) else: 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, momentum=BN_MOMENTUM), nn.ReLU(inplace=True))) transition_layers.append(nn.Sequential(*conv3x3s)) return nn.ModuleList(transition_layers)
HRnet是一個stage一個stage構建的，即橫向構建，每次都需要判斷pre_stage_channels和cur_stage_channels是否相等以判斷transition部分要做什么，這點非常好理解。
這是之前畫的可視化之后的結構圖，它的transition和論文中說的不同，但是我看了另外一個博主的文章她也是這么畫的所以估計是沒有問題的，好了終於把這篇博客寫完了，累啊，