人体姿态估计的深度高分辨率表示学习.
HRNet是一种基于热图的2D单人姿态估计模型，其核心思想是不断地去融合不同尺度上的信息。相比传统的下采样的网络结构，HRNet分成多个层级，但是始终保留着最精细的空间层级信息，通过融合下采样然后做上采样的层，来获得更多的上下文以及语义层面的信息（比如更大的感受野）。
1. HRNet的网络结构
HRNet的整体结构如图所示。网络整体由四个阶段构成。对于第$i=1,2,3,4$个阶段，网络并行地构造$2^2$-$2^{i+1}$倍下采样的特征图，并融合这些特征图以获得不同语义的信息。
HRNet首先通过两个卷积核大小为3x3、步长为2的卷积层共下采样了4倍。然后通过Layer1模块。Layer1是由重复堆叠的ResNet Bottleneck组成，只会调整通道个数，并不会改变特征层大小。
接着通过一系列Transition结构以及Stage结构，每通过一个Transition结构都会新增一个尺度分支。比如Transition1在layer1的输出基础上通过并行两个卷积核大小为3x3的卷积层得到两个不同的尺度分支，即下采样4倍的尺度以及下采样8倍的尺度。在Transition2中在原来的两个尺度分支基础上再新加一个下采样16倍的尺度。值得一提的是，每一个新的尺度分支都是通过一个卷积核大小为3x3、步长为2的卷积层构造的。
Stage结构用于融合不同尺度的特征。以Stage3为例，对于每个尺度分支，首先通过4个Basic Block，然后融合不同尺度上的信息。对于每个尺度分支上的输出都是由所有分支上的输出进行融合得到的。比如说对于下采样4倍分支的输出，它是分别将下采样4倍分支的输出（不做任何处理）、 下采样8倍分支的输出通过Up x2上采样2倍以及下采样16倍分支的输出通过Up x4上采样4倍进行相加，最后通过ReLU得到下采样4倍分支的融合输出。
对于所有的Up模块是通过一个卷积核大小为1x1的卷积层+Upsample直接放大n倍得到上采样后的结果（上采样默认采用最邻近插值）。Down模块每下采样2倍都要增加一个卷积核大小为3x3、步长为2的卷积层。
在Stage4中的最后一个Exchange Block只输出下采样4倍分支的输出（即只保留分辨率最高的特征层），然后接上一个卷积核大小为1x1卷卷积层。最终得到的特征层（64x48x17）就是针对每个关键点的热图。
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes, momentum=0.1)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes, momentum=0.1)
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=0.1)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes, momentum=0.1)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
momentum=0.1)
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 StageModule(nn.Module):
def __init__(self, input_branches, output_branches, c):
"""
构建对应stage，即用来融合不同尺度的实现
:param input_branches: 输入的分支数，每个分支对应一种尺度
:param output_branches: 输出的分支数
:param c: 输入的第一个分支通道数
"""
super().__init__()
self.input_branches = input_branches
self.output_branches = output_branches
self.branches = nn.ModuleList()
for i in range(self.input_branches): # 每个分支上都先通过4个BasicBlock
w = c * (2 ** i) # 对应第i个分支的通道数
branch = nn.Sequential(
BasicBlock(w, w),
BasicBlock(w, w),
BasicBlock(w, w),
BasicBlock(w, w)
)
self.branches.append(branch)
self.fuse_layers = nn.ModuleList() # 用于融合每个分支上的输出
for i in range(self.output_branches):
self.fuse_layers.append(nn.ModuleList())
for j in range(self.input_branches):
if i == j:
# 当输入、输出为同一个分支时不做任何处理
self.fuse_layers[-1].append(nn.Identity())
elif i < j:
# 当输入分支j大于输出分支i时(即输入分支下采样率大于输出分支下采样率)，
# 此时需要对输入分支j进行通道调整以及上采样，方便后续相加
self.fuse_layers[-1].append(
nn.Sequential(
nn.Conv2d(c * (2 ** j), c * (2 ** i), kernel_size=1, stride=1, bias=False),
nn.BatchNorm2d(c * (2 ** i), momentum=0.1),
nn.Upsample(scale_factor=2.0 ** (j - i), mode='nearest')
)
)
else: # i > j
# 当输入分支j小于输出分支i时(即输入分支下采样率小于输出分支下采样率)，
# 此时需要对输入分支j进行通道调整以及下采样，方便后续相加
# 注意，这里每次下采样2x都是通过一个3x3卷积层实现的，4x就是两个，8x就是三个，总共i-j个
ops = []
# 前i-j-1个卷积层不用变通道，只进行下采样
for k in range(i - j - 1):
ops.append(
nn.Sequential(
nn.Conv2d(c * (2 ** j), c * (2 ** j), kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(c * (2 ** j), momentum=0.1),
nn.ReLU(inplace=True)
)
)
# 最后一个卷积层不仅要调整通道，还要进行下采样
ops.append(
nn.Sequential(
nn.Conv2d(c * (2 ** j), c * (2 ** i), kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(c * (2 ** i), momentum=0.1)
)
)
self.fuse_layers[-1].append(nn.Sequential(*ops))
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
# 每个分支通过对应的block
x = [branch(xi) for branch, xi in zip(self.branches, x)]
# 接着融合不同尺寸信息
x_fused = []
for i in range(len(self.fuse_layers)):
x_fused.append(
self.relu(
sum([self.fuse_layers[i][j](x[j]) for j in range(len(self.branches))])
)
)
return x_fused
class HighResolutionNet(nn.Module):
def __init__(self, base_channel: int = 32, num_joints: int = 17):
super().__init__()
# Stem
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64, momentum=0.1)
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(64, momentum=0.1)
self.relu = nn.ReLU(inplace=True)
# Stage1
downsample = nn.Sequential(
nn.Conv2d(64, 256, kernel_size=1, stride=1, bias=False),
nn.BatchNorm2d(256, momentum=0.1)
)
self.layer1 = nn.Sequential(
Bottleneck(64, 64, downsample=downsample),
Bottleneck(256, 64),
Bottleneck(256, 64),
Bottleneck(256, 64)
)
self.transition1 = nn.ModuleList([
nn.Sequential(
nn.Conv2d(256, base_channel, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(base_channel, momentum=0.1),
nn.ReLU(inplace=True)
),
nn.Sequential(
nn.Sequential(
nn.Conv2d(256, base_channel * 2, kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(base_channel * 2, momentum=0.1),
nn.ReLU(inplace=True)
)
)
])
# Stage2
self.stage2 = nn.Sequential(
StageModule(input_branches=2, output_branches=2, c=base_channel)
)
# transition2
self.transition2 = nn.ModuleList([
nn.Identity(), # None, - Used in place of "None" because it is callable
nn.Identity(), # None, - Used in place of "None" because it is callable
nn.Sequential(
nn.Sequential(
nn.Conv2d(base_channel * 2, base_channel * 4, kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(base_channel * 4, momentum=0.1),
nn.ReLU(inplace=True)
)
)
])
# Stage3
self.stage3 = nn.Sequential(
StageModule(input_branches=3, output_branches=3, c=base_channel),
StageModule(input_branches=3, output_branches=3, c=base_channel),
StageModule(input_branches=3, output_branches=3, c=base_channel),
StageModule(input_branches=3, output_branches=3, c=base_channel)
)
# transition3
self.transition3 = nn.ModuleList([
nn.Identity(), # None, - Used in place of "None" because it is callable
nn.Identity(), # None, - Used in place of "None" because it is callable
nn.Identity(), # None, - Used in place of "None" because it is callable
nn.Sequential(
nn.Sequential(
nn.Conv2d(base_channel * 4, base_channel * 8, kernel_size=3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(base_channel * 8, momentum=0.1),
nn.ReLU(inplace=True)
)
)
])
# Stage4
# 注意，最后一个StageModule只输出分辨率最高的特征层
self.stage4 = nn.Sequential(
StageModule(input_branches=4, output_branches=4, c=base_channel),
StageModule(input_branches=4, output_branches=4, c=base_channel),
StageModule(input_branches=4, output_branches=1, c=base_channel)
)
# Final layer
self.final_layer = nn.Conv2d(base_channel, num_joints, kernel_size=1, stride=1)
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)
x = [trans(x) for trans in self.transition1] # Since now, x is a list
x = self.stage2(x)
x = [
self.transition2[0](x[0]),
self.transition2[1](x[1]),
self.transition2[2](x[-1])
] # New branch derives from the "upper" branch only
x = self.stage3(x)
x = [
self.transition3[0](x[0]),
self.transition3[1](x[1]),
self.transition3[2](x[2]),
self.transition3[3](x[-1]),
] # New branch derives from the "upper" branch only
x = self.stage4(x)
x = self.final_layer(x[0])
return x
2. HRNet的实现细节
网络最终输出的热图分辨率是原图的$\frac{1}{4}$，对每个关键点对应的预测信息求最大值的位置，即预测得分最大的位置，作为预测关键点的位置，映射回原图就能得到原图上关键点的坐标。
在实践时对于每个关键点并不是直接取得分最大的位置，而是分别对比该点左右两侧（x方向）、上下两侧（y方向）的得分，并对坐标进行偏移修正。
for n in range(coords.shape[0]):
for p in range(coords.shape[1]):
hm = batch_heatmaps[n][p]
px = int(math.floor(coords[n][p][0] + 0.5))
py = int(math.floor(coords[n][p][1] + 0.5))
if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
diff = np.array(
[
hm[py][px+1] - hm[py][px-1],
hm[py+1][px] - hm[py-1][px]
]
)
coords[n][p] += np.sign(diff) * .25
计算每个关键点对应的均方误差损失后，在计算总损失时并不是直接把每个关键点的损失进行相加，而是在相加前对于每个点的损失分别乘上不同的权重。下面给出了每个关键点的名称以及所对应的权重。
"kps": ["nose","left_eye","right_eye","left_ear","right_ear","left_shoulder","right_shoulder","left_elbow","right_elbow","left_wrist","right_wrist","left_hip","right_hip","left_knee","right_knee","left_ankle","right_ankle"]
"kps_weights": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5]
