HRNet论文及部分代码解释
一、HRNet论文阅读
原文：Deep High-Resolution Representation Learning for Human Pose Estimation
代码：pytorch-deep-high-resolution-net
论文关注的问题
明确论文想解决是姿态估计问题（Pose Estimation）。
对姿态估计问题的相关解释可参考：人体姿态估计的过去，现在，未来
论文介绍了一种以学习高分辨率特征的人体姿态估算的方法，称为HRNet（High-Resolution Net）。现在存在的大多数方法，都是提取图象的低分辨率特征，然后恢复成高分辨率特征进行预测。而HRNet的整个处理过程都以高分辨率为主，将高分辨子网作为第一阶段，逐步增加高分辨率到低分辨率的子网，同时将多分辨率子网并行连接。
新结构HRNet的提出
主体结构如下图
HRNet由多个平行，分辨率逐渐降低的子网络构成，平行的子网路通过多尺度融合进行信息交流。水平方向上为网络的深度，垂直方向上为特征图（分辨率）大小变化。
从图中可以看出，HRNet从始至终都有一条高分辨率子网连接（红色框内），并以此为主体逐步加入低分辨率的子网，同时并行的子网之间不断进行信息交流，而其他多数网络结构是从通过高分辨-低分辨率网络提取特征然后通过高分辨率到低分辨率网络产生的低分辨率表示来恢复高分辨率，具体做法可以是上采样、反卷积等。
这种结构的好处有以下两点：
- HRNet分辨率的降低，是使用新增平行子网的方式，而不是一连串的去降低特征分辨率，可以一直保持高分辨率特征，不需要有分辨率由低到高的处理过程，因此预测出来的关键点热图在空间上的分布会更加的精准；
- 与现在的大多数特征融合方式（使用深层特征图和浅层特征图进行融合）不同，HRNet使用相同深度和相似级别的低分辨率图像，通过重复的多尺度融合来提升高分辨率图像。因此，这种方式预测的热图可能更准确。
HRNet的实现
具体实现有以下几部分：
1. Sequential multi-resolution subnetworks
与当前主流的模型相似，HRNet也有一个high-low分辨率的子网连接，称为连续多分辨率子网，每个子网形成一个stage，由一系列卷积组成，并且在相邻的子网之间有一个下采样层来将分辨率减半。
给定子网编号
N
s
r
N_{sr}
Nsr，下标
s
s
s代表第s个stage的子网，下标
r
r
r为分辨率序号，则子网
N
s
r
N_{sr}
Nsr的分辨率为第一个子网分辨率的
1
2
r
−
1
\frac{1}{2^{r-1}}
2r−11，这里的high-to-low网络结构如下
2. Parallel multi-resolution subnetworks
上面提到，HRNet用一个高分辨率的子网作为第一阶段，逐渐添加由高到低的子网层组成新的stage，并将这些平行多分辨率的子网模块连接起来，后一阶段并行子网的分辨率由前一阶段的分辨率和一个更低的分辨率组成。例如一个包含4个并行子网的网络结构如下图所示
3. Repeated multi-scale fusion
HRNet引入了平行网络信息交换单元，每个子网络重复接受来自其他平行子网络的信息。下面是一个例子，展示了信息交换的方案。将第三阶段分成了3个交换模块，每个模块由3个并行卷积单元和一个跨并行单元的交换单元组成，其结构如下：
其中
C
s
r
b
C^b_{sr}
Csrb表示第s个阶段的第b个模块中第r个分辨率序号的卷积单元，
ε
s
b
\varepsilon^b_s
εsb为相应的信息交换单元，不同单元之间交换信息的方式如下图
上图从左到右依次展示了交换单元如何聚合高、中、低分辨率的信息，最右边的图示 strided 3×3 表示使用 3x3 的跨步卷积，up samp. 1×1表示使用1x1的卷积进行上采样。
现在仅关注每一个阶段的输入和输出，假定输入为
{
X
1
,
X
2
,
.
.
.
,
X
s
}
\{X_1, X_2, ..., X_s\}
{X1,X2,...,Xs}，输出为
{
Y
1
,
Y
2
,
.
.
.
,
Y
s
}
\{Y_1, Y_2, ... , Y_s\}
{Y1,Y2,...,Ys}，输入和输出的特征大小是一致的，每个输出都是输入集合的映射，计算方式为
Y
k
=
Σ
i
=
1
s
a
(
X
i
,
k
)
Y_k=\Sigma_{i=1}^sa(X_i,k)
Yk=Σi=1sa(Xi,k)对于跨阶段的交换单元，会有一个额外的输出映射
Y
s
+
1
=
a
(
X
s
,
s
+
1
)
Y_{s+1}=a(X_s,s+1)
Ys+1=a(Xs,s+1)。
函数 a ( X i , s ) a(X_i,s) a(Xi,s)包含了分辨率 i i i到 k k k的上采样或者下采样过程，在这里，使用步长为2的 3×3 卷积进行2倍降维，使用简单的最邻近抽样方法进行上采样，而 1×1 卷积则用于对齐通道数。如果 i = k , a ( . , . ) i=k,a(., .) i=k,a(.,.)是一个恒等连接，则 a ( X i , k ) = X i a(X_i, k)=X_i a(Xi,k)=Xi。
4. Heatmap estimation
HRNet通过最后一个交换单元的高分辨率输出来回归获得热图，使用MSE损失函数比较预测热图和标注热图。
5. Network instantiation
HRNet由四个阶段组成，包含四个平行的子网络，在每个模块中，若分辨率降低到一半，则其通道数增加到原来的两倍。
第一个阶段包含四个残差模块，其结构与ResNet-50相同，是由一个通道为64的bottleneck形成，然后进行一个 3×3 的卷积，将特征图的通道数减小到 C C C，第2、3、4阶段共有1、4、3个信息交换模块，一个信息交换模块包含4个剩余单元，其中每个单元在每个分辨率中包含两个 3×3 卷积以及跨分辨率的交换单元。所以，共有8个交换单位，即进行了8次多尺度融合。
作者已经实现了两个不同大小的HRNet模型，分别为HRNet-W32和HRNet-W48，32和48代表高分辨率子网在最后三个阶段的宽度 C C C，其他三个并行子网的宽度为：HRNet-W32 = 64, 128, 256，HRNet-W48 = 96, 192, 384，并且使用256×192和384×288两种不同分辨率输入进行实验。
其他结构
在论文中，作者将HRNet模型在COCO2017和MPII两个数据集上分别测试，并与其他state-of-the-art模型进行比较，如Hourglass（Hourglass Networks）、CPN（Cascaded Pyramid Network）和SimpleBaseline，它们的结构分别如下图
(
a
)
(
b
)
(
c
)
(a)(b)(c)
(a)(b)(c)所示
对于具体结构的解析可参考以下几篇文章：
- Stacked Hourglass Networks for Human Pose Estimation(Hourglass)
- Cascaded Pyramid Network for Multi-Person Pose Estimation(CPN)
- Simple Baselines for Human Pose Estimation and Tracking
与这些结构比较，HRNet有两大方面的优势，使得其具有更好的效果：
- 在整个过程中保持高分辨率，而不需要从低分辨率恢复高分辨率表示
- 反复融合多分辨率表示，以呈现可靠的高分辨率表示
二、对应代码标注
训练文件（/tools/train.py）
代码及相关注释
def parse_args():
parser = argparse.ArgumentParser(description='Train keypoints network')
# general
# 加载指定的yaml配置文件
parser.add_argument('--cfg',
help='experiment configure file name',
required=True,
type=str)
parser.add_argument('opts',
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
# philly
# 模型目录
parser.add_argument('--modelDir',
help='model directory',
type=str,
default='')
# 日志目录
parser.add_argument('--logDir',
help='log directory',
type=str,
default='')
# 训练数据目录
parser.add_argument('--dataDir',
help='data directory',
type=str,
default='')
# 预训练模型目录
parser.add_argument('--prevModelDir',
help='prev Model directory',
type=str,
default='')
args = parser.parse_args()
return args
def main():
# 解析输入参数，并更新cfg
args = parse_args()
update_config(cfg, args)
# 配置日志记录
logger, final_output_dir, tb_log_dir = create_logger(
cfg, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(cfg)
# cudnn related setting
# GPU设置
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
# 使用配置文件构建网络
model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
cfg, is_train=True
)
# copy model file
this_dir = os.path.dirname(__file__)
shutil.copy2(
os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'),
final_output_dir)
# logger.info(pprint.pformat(model))
# 用于训练信息的图形化显示
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# 用于模型的图形化显示
dump_input = torch.rand(
(1, 3, cfg.MODEL.IMAGE_SIZE[1], cfg.MODEL.IMAGE_SIZE[0])
)
writer_dict['writer'].add_graph(model, (dump_input, ))
logger.info(get_model_summary(model, dump_input))
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
# define loss function (criterion) and optimizer
criterion = JointsMSELoss(
use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT
).cuda()
# Data loading code
normalize = transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
)
# 创建训练、测试数据集及其迭代器
train_dataset = eval('dataset.'+cfg.DATASET.DATASET)(
cfg, cfg.DATASET.ROOT, cfg.DATASET.TRAIN_SET, True,
transforms.Compose([
transforms.ToTensor(),
normalize,
])
)
valid_dataset = eval('dataset.'+cfg.DATASET.DATASET)(
cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
transforms.Compose([
transforms.ToTensor(),
normalize,
])
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
shuffle=cfg.TRAIN.SHUFFLE,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY
)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
shuffle=False,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY
)
# 模型加载以及优化策略的相关配置
best_perf = 0.0
best_model = False
last_epoch = -1
# 获取优化器配置
optimizer = get_optimizer(cfg, model)
begin_epoch = cfg.TRAIN.BEGIN_EPOCH
checkpoint_file = os.path.join(
final_output_dir, 'checkpoint.pth'
)
# 断点的自动恢复
if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
begin_epoch = checkpoint['epoch']
best_perf = checkpoint['perf']
last_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_file, checkpoint['epoch']))
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,
last_epoch=last_epoch
)
# 开始训练
for epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH):
lr_scheduler.step()
# train for one epoch
train(cfg, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# evaluate on validation set
perf_indicator = validate(
cfg, valid_loader, valid_dataset, model, criterion,
final_output_dir, tb_log_dir, writer_dict
)
# 对比结果，保存最优模型
if perf_indicator >= best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
# 保存断点
save_checkpoint({
'epoch': epoch + 1,
'model': cfg.MODEL.NAME,
'state_dict': model.state_dict(),
'best_state_dict': model.module.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
final_model_state_file = os.path.join(
final_output_dir, 'final_state.pth'
)
logger.info('=> saving final model state to {}'.format(
final_model_state_file)
)
torch.save(model.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
if __name__ == '__main__':
main()
train代码流程
1. 简单流程
2. 详细流程
基本模块
conv3x3
返回一个核为3×3，步长为1，padding为1的2d卷积层
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding 1"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
BasicBlock
定义基础的残差模块，结构类似于下图
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
Bottleneck模块
ResNet中的Bottleneck模块，结构与下图类似
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)
# 输入维度为planes，输出维度为 planes*4
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
HighResolutionModule高分辨率模块
在该模块中进行平行子网络的信息交流
__init__函数
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
"""
:param num_branches: 当前stage分支的平行子网络数目
:param blocks: BasicBlock或者Bottleneck
:param num_blocks: BasicBlock或者Bottleneck数目
:param num_inchannels: 输入通道数
stage=2时, num_inchannels=[32,64]
stage=3时, num_inchannels=[32,64,128]
stage=4时, num_inchannels=[32,64,128]
:param num_channels: 输出通道数
stage=2时, num_channels=[32,64]
stage=3时, num_channels=[32,64,128]
stage=4时, num_channels=[32,64,128,256]
:param fuse_method: 融合方式，默认为fuse
:param multi_scale_output: HRNet-W32配置下，所有stage均为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
# 为每个分支构建分支网络
# stage=2,3,4，平行子网分支数num_branches=2,3,4
self.branches = self._make_branches(
num_branches, blocks, num_blocks, num_channels)
# 创建多尺度融合层
# stage=2,3,4时，len(self.fuse_layers)分别为2,3,4，与其num_branches在每个stage的数目是一致的
self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(True)
_check_branches函数
对输入的一些参数进行检测，判断num_branches 和 num_blocks, num_inchannels, num_channels (list) 三者的长度是否一致，不一致则报错
_make_one_branch函数
该函数作用为搭建一个分支，由num_branches[branch_index]个BasicBlock或Bottleneck组成。
def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
# 如果stride不为1, 或者输入通道数目与输出通道数目不一致，使用1×1卷积改变通道数
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 = []
# 为当前分支branch_index创建一个block,在该处进行下采样
layers.append(
block(
self.num_inchannels[branch_index],
num_channels[branch_index],
stride,
downsample
)
)
# 将输出通道数赋值给输入通道数,为下一stage作准备
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
# # 为[1, num_blocks[branch_index]]分支创建block
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)
_make_branches函数
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)
_make_fuse_layers
作用：进行多个分辨率层的融合。
def _make_fuse_layers(self):
if self.num_branches == 1:
return None
# 平行子网分支数
num_branches = self.num_branches
# 输入通道数
num_inchannels = self.num_inchannels
fuse_layers = []
# 为每个分支都创建对应的特征融合网络,如果multi_scale_output=1,则只需要一个特征融合网络
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
# 多种情况的讨论
# 1.当前分支信息传递到上一分支(沿论文图示scale方向)的下一层(沿论文图示depth方向),进行上采样,分辨率加倍
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.当前分支信息传递到当前分支(论文图示沿scale方向)的下一层(沿论文图示depth方向),不做任何操作,分辨率相同
elif j == i:
fuse_layer.append(None)
# 3.当前分支信息传递到当前分支(论文图示沿scale方向)的下一层(沿论文图示depth方向),分辨率减半,分辨率减半
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)
)
)
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)
)
)
fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
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
主模型PoseHighResolutionNet代码标注
这里以HRNet-W32，输入大小 256×192 的模型为例，HRNet主体模块的生成与对应维度变化如下图所示：
注：原文中对第一阶段的描述
The first stage contains 4 residual units where each unit, the same to the ResNet-50, is formed by a bottleneck with the width 64, and is followed by one 3×3 convolution reducing the width of feature maps to C.
第一阶段采用了ResNet-50结构，包含了四个Bottleneck，最后输出维度为256，所以 N 11 N_{11} N11维度为[b, 256, 64, 48]。而其他阶段的通道数，在源码的配置文件w32_256x192_adam_lr1e-3.yaml中已经给出，如下图对stage2和stage3的相关配置：
部分代码标注如下：
__init__函数
首先是PoseHighResolutionNet的__init__函数，主要对网络各个层及相关参数进行定义
def __init__(self, cfg, **kwargs):
# 初始输入维度为64
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)
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']] # 在配置文件中为BASIC
# 配置stage几个模块的通道数
num_channels = [
# BASIC的expansion为 1，所以stage2的第一个分支为通道数32，第二个为64
num_channels[i] * block.expansion for i in range(len(num_channels)) # len(num_channels)=2
]
# 生成新的平行分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)
# 与stage2过程相似，其维度计算方式均可通过配置文件计算获取，不再详细标注
self.stage3_cfg = extra['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition2 = self._make_transition_layer(
pre_stage_channels, num_channels)
self.stage3, pre_stage_channels = self._make_stage(
self.stage3_cfg, num_channels)
# 与stage2过程相似，其维度计算方式均可通过配置文件计算获取，不再详细标注
self.stage4_cfg = extra['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition3 = self._make_transition_layer(
pre_stage_channels, num_channels)
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']
_make_transition_layer函数
在初始化函数中，生成新分支主要使用_make_transition_layer函数，该函数的作用在于生成下一阶段同等分辨率和一般分辨率的分支。首先会进行循环遍历，对每个分支进行处理，详细的处理方式在代码中标注。
def _make_transition_layer(
self, num_channels_pre_layer, num_channels_cur_layer):
"""
:param num_channels_pre_layer: 上一个stage平行网络的输出通道数目,为一个list,
stage=2时, num_channels_pre_layer=[256]
stage=3时, num_channels_pre_layer=[32,64]
stage=4时, num_channels_pre_layer=[32,64,128]
:param num_channels_cur_layer:
stage=2时, num_channels_cur_layer = [32,64]
stage=3时, num_channels_cur_layer = [32,64,128]
stage=4时, num_channels_cur_layer = [32,64,128,256]
"""
num_branches_cur = len(num_channels_cur_layer)
num_branches_pre = len(num_channels_pre_layer)
transition_layers = []
# 对stage的每个分支进行处理
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]),
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),
nn.ReLU(inplace=True)
)
)
transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
_make_stage函数
该函数的作用在于生成stage时的HighResolutionModule。
def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
"""
当stage=2时： num_inchannels=[32,64] multi_scale_output=Ture
当stage=3时： num_inchannels=[32,64,128] multi_scale_output=Ture
当stage=4时： num_inchannels=[32,64,128,256] multi_scale_output=False
"""
# 当stage=2,3,4时,num_modules分别为：1,4,3
# 表示HighResolutionModule（平行之网络交换信息模块）模块的数目
num_modules = layer_config['NUM_MODULES']
# 当stage=2,3,4时,num_branches分别为：2,3,4,表示每个stage平行网络的数目
num_branches = layer_config['NUM_BRANCHES']
# 当stage=2,3,4时,num_blocks分别为：[4,4], [4,4,4], [4,4,4,4],
# 表示每个stage blocks(BasicBlock或者BasicBlock)的数目
num_blocks = layer_config['NUM_BLOCKS']
# 当stage=2,3,4时,num_channels分别为：[32,64],[32,64,128],[32,64,128,256]
# 在对应stage, 对应每个平行子网络的输出通道数
num_channels = layer_config['NUM_CHANNELS']
# 根据配置文件，HRNet-W32的所有stage的block均为BasicBlock
block = blocks_dict[layer_config['BLOCK']]
# 根据配置文件，HRNet-W32的所有stage的FUSE_METHOD(融合方式)均为SUM
fuse_method = layer_config['FUSE_METHOD']
modules = []
# 根据num_modules的数目创建HighResolutionModule
for i in range(num_modules):
# multi_scale_output is only used last module
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
# # 根据参数,添加HighResolutionModule
modules.append(
HighResolutionModule(
num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output
)
)
# 获得最后一个HighResolutionModule的输出通道数
num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
forward函数
forward函数的作用已不必再介绍，直接看代码了解处理过程即可
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
参考：
2573
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
