HRNet网络代码解读：Deep High-Resolution Representation Learning for Human Pose Estimation

论文地址：https://arxiv.org/abs/1902.09212
官方源码：https://github.com/leoxiaobin/deep-high-resolution-net.pytorch

HRNet网络结构

这里引用“太阳花的小绿豆”绘制的一张基于HRNet-32模型的结构图。便于后续理解。

重要的部分写在代码注释里了，阅读的时候注意。

代码详解

forward函数

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类，让我们进入到这个类看一下。

首先看forward函数：

    def forward(self, x):
        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)

所对应的stem net为：

# stem net
        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)

这里经过两个卷积bn激活函数的操作，后接一个layer1模块，特征通道数下采样4倍，通道变为256.

其中layer1由_make_layer(Bottleneck, 64, 4)构建。

layer1函数

让我们看下_make_layer的具体操作。

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.inplanes, planes * block.expansion,
                    kernel_size=1, stride=stride, bias=False
                ),
                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        # 通道数由64变为256
        self.inplanes = planes * block.expansion
        # self.inplanes = 64 * 4 = 256
        for i in range(1, blocks):
        	# 重复堆叠三次，不使用downsample，其实这里的downsample操作也并没有进行下采样。
        	# 输入通道数为256，输出通道数也为256
        	# 最后得到特征图的大小为下采样4倍，输出通道256的featuremap
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

Bottleneck类中expansion = 4, self.inplanes = 64 != 64 *4 执行downsample操作。注意这里downsample并没有对模型进行下采样，stride= 1，只是沿用了Resnet的名称，叫成了downsample。

Bottleneck的搭建如下代码：

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self