姿态估计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] ---> y[b, 32, 64, 48] 因为通道数不一致,通过卷积进行通道数变换
# y[b, 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[b, 32, 64, 48] ---> y[b, 32, 64, 48]
# x[b, 64, 32, 24] ---> y[b, 64, 32, 24]
y_list = self.stage2(x_list)
# 对应论文中的stage3
# 其中包含了创建分支的过程,即 N22-->N32,N33 这个过程
# N33的分辨率为N32的二分之一,
# y[b, 32, 64, 48] ---> x[b, 32, 64, 48] 因为通道数一致,没有做任何操作
# y[b, 64, 32, 24] ---> x[b, 64, 32, 24] 因为通道数一致,没有做任何操作
# x[b, 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[b, 32, 64, 48] ---> x[b, 32, 64, 48]
# x[b, 32, 32, 24] ---> x[b, 32, 32, 24]
# x[b, 64, 16, 12] ---> x[b, 64, 16, 12]
y_list = self.stage3(x_list)
# 对应论文中的stage4
# 其中包含了创建分支的过程,即 N33-->N43,N44 这个过程
# N44的分辨率为N43的二分之一
# y[b, 32, 64, 48] ---> x[b, 32, 64, 48] 因为通道数一致,没有做任何操作
# y[b, 64, 32, 24] ---> x[b, 64, 32, 24] 因为通道数一致,没有做任何操作
# y[b, 128, 16, 12] ---> x[b, 128, 16, 12] 因为通道数一致,没有做任何操作
# x[b, 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[b, 32, 64, 48] --->
# x[b, 64, 32, 24] --->
# x[b, 128, 16, 12] --->
# x[b, 256, 8, 6 ] ---> y[b, 32, 64, 48]
y_list = self.stage4(x_list)
# y[b, 32, 64, 48] --> x[b, 17, 64, 48]
x = self.final_layer(y_list[0])
return x
其上为两个部分，分别为初始化过程,以及前向传播的过程。
forward
对于前线传播过程，大家看了注释之后应该是很清楚很明白了,主要对应论文中的如下过程：
N 11 → N 21 → N 31 → N 41 ↘ N 22 → N 32 → N 42 ↘ N 33 → N 43 ↘ N 44 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} N11→N21→N31→N41↘N22→N32→N42↘N33→N43↘N44
通过代码我们可以很明显的看到，最终我们获得一个[b, 17, 64, 48]大小的heatmap，这就是我们最终想要的结果，其上的每个 N X X N_{XX} NXX可以分成两个重要的模块，分别为 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是为了构建论文中平行子网络信息交流的模块。具体的实现过程,在下篇博客中为大家讲解。