PRTR论文代码解读
Pose Recognition with Cascade Transformers
paper:https://arxiv.org/abs/2104.06976
code:https://github.com/mlpc-ucsd/PRTR
这里对PRTR论文进行解读记录，方便自己以后学习。
1 HRNET网络结构解析
目前，深度卷积神经网络提供了主流的解决方案。主要有两种方法：回归关键点位置 regressing the position of keypoints 和估算关键点热图 estimating keypoint heatmaps，然后选择热值最高的位置作为关键点。
1.1 High-to-low and low-to-high
high-to-low 的目标是生成低分辨率和高分辨率的表征，low-to-high 的目标是生成高分辨率的表征。这两个过程可能会重复多次，以提高性能。增加多尺度信息之间的融合是非常有效的，例如原图像和模糊图像进行联合双边滤波可以得到介于两者之间的模糊程度的图像，而RGF滤波就是重复将联合双边滤波的结果作为那张模糊的引导图，这样得到的结果会越来越趋近于原图。同理，不同分辨率的图像采样到相同的尺度反复的融合，加之网络的学习能力，会使得多次融合后的结果更加趋近于正确的表示。
现有的网络设计模式有：
- 对称结构，先下采样，再上采样，同时使用跳层连接恢复下采样丢失的信息；
- 级联多字金字塔；
- 先下采样，转置卷积上采样，不使用跳层连接进行数据融合；
- 扩张卷积，减少下采样次数，不使用跳层连接进行数据融合；
如下图所示：
1.2 多分辨率子网络
并行高分辨率子网
以高分辨率子网为第一阶段，逐步增加高分辨率到低分辨率的子网，形成新的阶段，并将多分辨率子网并行连接。因此，后一阶段并行子网的分辨率由前一阶段的分辨率和下一阶段的分辨率组成。一个包含4个并行子网的网络结构示例如下：
重复多尺度融合
HRNet中引入了跨并行子网的交换单元，使每个子网重复接收来自其他并行子网的信息。下面是一个展示信息交换方案的示例。将第三阶段划分为若干个交换块（如3个），每个块由3个并行卷积单元与1个交换单元跨并行单元进行卷积，得到:
1.3 HRNet 结构
HRNet 主要的模型结构，具体实现部分在 HighResolutionNet 类中有详细定义。
总体结构 按照顺序 可分为三部分：
1. stem net：
- 从 IMG 到 1/4 大小的 feature map，得到此尺寸的特征图后，之后的 HRNet 始终保持此尺寸的图片
2. HRNet 4 stages：如下图所示的 4 阶段 由 HighResolutionModule 组成的模型
- 其中，每个蓝色底色为1个阶段
- 每个 stage 产生的 multi-scale 特征图，具体配置如下表，以 hrnet_48 为例
- stage 的连接处有 transition 结构，用于在不同 stage 之间连接，完成 channels 及 feature map 大小对应
| multi-scale feature map | num_branches (分支数) | num_blocks (每个分支 block 重复次数) | num_modules (HighResolutionModule 重复次数) | |
| stage1 | [1/4] | 1 | [4] | 0 |
| stage2 | [1/4, 1/8] | 2 | [4，4] | 1 |
| stage3 | [1/4, 1/8, 1/16] | 3 | [4，4，4] | 4 |
| stage4 | [1/4, 1/8, 1/16, 1/32] | 4 | [4，4，4，4] | 3 |
3. segment head：
- 将 stage4 输出的 4 种 scale 特征 concat 到一起
- 加上
num_channels -> num_classes层，得到分割结果
1.4 HRNet 构建函数 def HRNet(cfg_path, **kwargs)
- 通过指定
cfg_path选择要使用的模型的结构（yaml 存储） - 通过指定 kwargs 选择是否选用 pretrain 模型
具备 pretrain 模型的，可用模型结构：
- seg_hrnet_w18_small_v2_sgd_lr5e-2_wd1e-4_bs32_x100.yaml
- seg_hrnet_w30_sgd_lr5e-2_wd1e-4_bs32_x100.yaml
- seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml，为目前采用的结构
def HRNet(cfg_path, **kwargs):
from models.hrnet.config import update_config
cfg = update_config(cfg_path)
model = HighResolutionNet(cfg, **kwargs)
if kwargs.get('use_pretrain', False):
model.load_pretrain(cfg.MODEL.PRETRAINED)
return modelyaml 文件中，关于模型结构的关键部分，以 hrnet_w48 为例
MODEL:
NAME: seg_hrnet
ALIGN_CORNERS: True
PRETRAINED: 'pretrained_models/hrnetv2_w48_imagenet_pretrained.pth' # 指定 pretrain 模型路径
EXTRA: # EXTRA 具体定义了模型的结果，包括 4 个 STAGE，各自的参数
FINAL_CONV_KERNEL: 1
STAGE1:
NUM_MODULES: 1
NUM_RANCHES: 1
BLOCK: BOTTLENECK
NUM_BLOCKS:
- 4
NUM_CHANNELS:
- 64
FUSE_METHOD: SUM
STAGE2:
NUM_MODULES: 1 # HighResolutionModule 重复次数
NUM_BRANCHES: 2 # 分支数
BLOCK: BASIC
NUM_BLOCKS:
- 4
- 4
NUM_CHANNELS:
- 48
- 96
FUSE_METHOD: SUM
STAGE3:
NUM_MODULES: 4
NUM_BRANCHES: 3
BLOCK: BASIC
NUM_BLOCKS:
- 4
- 4
- 4
NUM_CHANNELS:
- 48
- 96
- 192
FUSE_METHOD: SUM
STAGE4:
NUM_MODULES: 3
NUM_BRANCHES: 4
BLOCK: BASIC
NUM_BLOCKS:
- 4
- 4
- 4
- 4
NUM_CHANNELS:
- 48
- 96
- 192
- 384
FUSE_METHOD: SUM2 源码结构
下表列出HRNet中比较重要的文件：
| 文件名称 | 功能 |
|---|---|
| tools/trian.py | 训练脚本 |
| tools/test.py | 测试脚本 |
| lib/dataset/mpii.py | 对MPII数据集进行预处理 |
| lib/dataset/JointsDataSet | 数据读取脚本 |
| lib/models/pose_hrnet.py | 网络结构构建脚本 |
| lib/utils | HRNet的一些方法 |
| experiments/mpii/hrnet | HRNet网络的初始化参数脚本 |
接下来对一些重要文件，将一一讲解，并且说清数据流的走向和函数调用关系。
代码的总体结构如上图所示
1. data中有coco的ann，images，后一个person_detect_result是MS自己测试出来的框图结果。
2. experiment是网络训练中保存的参数，一般以yaml格式进行存储。针对不同的resnet，设置了不同的超参数数值。
3. lib内包含所有工程代码，
- core中包含config，evaluate，function，inference，loss四个函数。
- dataset包含继承nn.DataSet的JointsDataset，用于实现getitem方法，coco和mpii为从不同数据集获取图像的方法。用于解耦。
- models包含pose_resnet，为模型的核心代码，继承了resnet非全连接层写法，并在后方添加了反卷积直接输出。nms方法，用于提升结果精度。
- utils方法中包含transform，utils.py->create_logger, get_optimizer, save_checkpoint，vis保存图像，zipreader.
4. log记录按时间打的代码
5. models按不同的数据集，不同的backbone和inputsize记录了不同的ckpt
6. otuput按数据集->backbone->inputsize打log
7. pose_estimation包含train和valid两个算法，train的过程也执行了valid算法。
代码的总体逻辑：
使用JointDataSet进行获取数据后的处理，从coco数据集获取的图像仅仅是一个左上右下的框，如何将原始的框整合为统一尺度输入网络？其策略是：首先，当然不能resize，这样形状转变容易影响对于一个人的判断。它将尺度不一致的框，若w:h不是256：192,那么将w或者h进行扩充，直到框的比例达到256：192。其次，将图像的w，h与200做比，计算出图像的scale。第三步，已经确定了一个点center，center的位置就是进行仿射变换的256*192的图像中点。框图左上角的坐标就是输入图像的0，0坐标点，两点再确定一个中点，得到三个点就可以进行仿射变换了。图像仿射变换后需要记录原始图像的center和scale，以便之后逆变换回来。
图像输入网络时，网络使用的结构为resnet+反卷积网络结构。论文中说明，当heatmap为原图的1/4的时候，acc最高，所以没有恢复到原图的整体大小直接输出。输出之后将图像进行反转后再检测一次，融合两次的检测结果达到最高的精度。该网络使用常规的heatmap输出网络的结果，输出17张特征图，每张特征图的weight最大的时候，就是关键点的位置。使用debug模式将图像映射到torchvision.grids上，然后将grid转为numpy，将每个坐标乘以4之后(因为输出的heatmap是原始输入图片的1/4)，显示到图片上。
使用dataset.evaluate将图像的结果还原到整张大图上。该方法为，已知output的尺寸和output上的各个点，并且知道center和scale，也就是框图在原始图像中的位置，那么可以利用center和scale计算出三个点，仿射变换到output上。求出仿射矩阵。但是这时候需要反向映射到原图上，所以输入的三个点是output的三个点，映射的三个点是center上的三个点，也就相当于是一个逆仿射变换。
然后使用nms算法重定向？这一块还需要仔细查看一下具体做法，然后映射得坐标和框和图片整合成一个json保存下来作为最终得结果。
感受：第一次接触关键点检测感觉很神奇，首先关键点得映射我一直认为是一个回归问题，但是实际做法项目做成了分类问题。而且本文是个up-bottom得方法，先检测出框，然后将整个框缩放之后进行训练，也就解释了多人效果为啥那么好。还有就是映射得问题我一直没有看懂，为啥center与scale就能确定一个框，scale不是尺度吗？其实这个scale是基于200得比值结果，是个框得大小的意思。另外输出结果后，在debug模式我都没有搞明白为什么pred*4，跑去看model，看了半天才发现前面就是个纯resnet101，输出为7*7，后面为反卷积，反卷积了三次，也就是热图是原图的1/4，看论文才意识到热图为原图1/4时效果最好。所以pred*4之后的点才是输入图的关键点位置。最后一个evalute方法中output映射到原图的时候我一直在纠结output是原图的1/4，映射到output上不应该结果的关键点应该变为原图的1/4了吗？后来发现output的尺寸是输入，关键点的映射目标位置为center和scale，想反了。
3 get_pose_net搭建模型
文件路径：\PRTR-main\two_stage\lib\models\pose_transformer.py
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这里面用的model=PoseHighResolutionNet(cfg, **kwargs)来构建整个模型，所以我们来看PoseHighResolutionNet(cfg, **kwargs)类的forward函数，并且一节一节开始分析。
3.1 初步提特征
首先是最简单的一节，这一节就是先对输入的图片进行简单的提取特征，没啥好说的，自己对照这init函数看看就晓得了。
def forward(self, x):
#初步的进行提取特征
x = self.conv1(x) #(h,w,3)-->((hin+1)/2,(win+1)/2,64)
x = self.bn1(x) #正则化
x = self.relu(x) #激活函数
x = self.conv2(x) #(h,w,64)-->((hin+1)/2,(win+1)/2,64)
x = self.bn2(x) #正则化
x = self.relu(x) #激活函数模型结构是这样的：
(conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)3.2 利用残差结构，加深层数继续提特征
在forward函数中，初步提特征后下一行是：
x = self.layer1(x)我们先来看看self.layer1在init中的定义：
self.layer1 = self._make_layer(Bottleneck, 64, 4)然后我们再进入到self._make_layer(Bottleneck, 64, 4)函数去看
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
#我们来看一下下面的if部分
#在layer1中，block传入的是Bottlenect类，block.expansion是block类里的一个变量，定义为4
#layer1的stride为1，planes为64，而self.inplane表示当前特征图通道数，经过初步提特征处理后的特征图通道数为是64，block.expanson=4，达成条件
#那么downsample = nn.Sequential(
# nn.Conv2d(64, 64*4,kernel_size=1, stride=1, bias=False),
# nn.BatchNorm2d(64*4, momentum=BN_MOMENTUM),
# )
#这里的downsample会在后面的bottleneck里面用到，用于下面block中调整输入x的通道数，实现残差结构相加
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里第一层是：bottleneck(64, 64, 1, downsample) (w,h,64)-->(w,h,256) 详细的分析在下面哦
layers.append(block(self.inplanes, planes, stride, downsample))
#经过第一层后，当前特征图通道数为256
self.inplanes = planes * block.expansion
#这里的block为4，即for i in range(1,4)
#所以这里for循环实现了3层bottleneck，目的应该是为了加深层数
#bottleneck(256, 64, 1) 这里就没有传downsample了哦，因为残差结构相加不需要升维或者降维
#bottleneck(256, 64, 1)
#bottleneck(256, 64, 1)
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)#############################################################################
以layer1的第一层bottleneck(64, 64, 1, downsample)为例子，我们再来看看bottleneck到底干了个啥，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)
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) #n.Conv2d(64,64, kernel_size=1, bias=False) （w,h,64)-->(w,h,64)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out) #nn.Conv2d(64, 64, kernel_size=3, 1,padding=1, bias=False) (w,h,64)-->(w,h,64)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out) #nn.Conv2d(64, 64 * 4, kernel_size=1,bias=False) (w,h,64)-->(w,h,256)
out = self.bn3(out)
if self.downsample is not None:
#这里的downsample的作用是希望输入原图x与conv3输出的图维度相同，方便两种特征图进行相加，保留更多的信息（你要是看不懂这句话，就去先简单了解一下残差结构）
#如果x与conv3输出图维度本来就相同，就意味着可以直接相加，那么downsample会为空，自然就不会进行下面操作
residual = self.downsample(x) #downsample = nn.Sequential(
# nn.Conv2d(64, 64*4,kernel_size=1, stride=1, bias=False),
# nn.BatchNorm2d(64*4, momentum=BN_MOMENTUM),
#
out += residual #残差结构相加嘛
out = self.relu(out) 得到结果
return out#############################################################################
那么这一部分的模型结构是这样子滴
(layer1): Sequential(
(0): Bottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
)
)那么我们就完成了总体特征图的这个部分
3.3 添加分支_make_transition_layer
接着往下看forward代码
x_list = []
#我们先看这个循环条件，在配置文件中self.stage2_cfg['NUM_BRANCHES']为2（其实总结构图上不也是画着两个分支嘛，分支也可以理解为有多少份不同尺寸的特征图）
#所以这里有两个循环，i=0或1
#在init中，有几行代码与self.transition1[i]有关，我们先搞清楚self.transition1[i]里到底是啥
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)
y_list = self.stage2(x_list)在init中与self.transition1[i]有关的代码块：
''''''
extra['STAGE2']为
STAGE2:
NUM_MODULES: 1
NUM_BRANCHES: 2
BLOCK: BASIC
NUM_BLOCKS:
- 4
- 4
NUM_CHANNELS:
- 32
- 64
FUSE_METHOD: SUM
''''''
self.stage2_cfg = extra['STAGE2']
#num_channels此时为[32,64]，
num_channels = self.stage2_cfg['NUM_CHANNELS']
#block为basic，传入的是一个类BasicBlock，因为代码中定义了一个blocks_dict = {'BASIC': BasicBlock,'BOTTLENECK': Bottleneck}
block = blocks_dict[self.stage2_cfg['BLOCK']]
#num_channels =[32*1,64*1]，这里num_channels的意义是stage2中，各个分支的通道数，这里乘1是因为basicblock里面expansion是1，即残差结构不会扩展通道数
num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))]
#这里有引入一个新的函数self._make_transition_layer
self.transition1 = self._make_transition_layer([256], num_channels)于是我们再看看self._make_transition_layer这个函数到底做了什么
文件路径：\PRTR-main\two_stage\lib\models\hrnet.py
#两个参数，num_channels_pre_layer是之前每个分支的通道数，stage1的时候只有一个分支，通道数为256
#num_channels_cur_layer完成transition之后每个分支的通道数，这个上面已经设置好了，在stage1的时候为[32,64]
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 = []
#stage1的时候，num_branches_cur为2，所以有两个循环，i=0、1
for i in range(num_branches_cur):
# 由于branches_cur有两个分支，branches_pre只有一个分支，
#所以我们可以直接利用branches_pre已有分支作为branches_cur的其中一个分支
#这个操作是hrnet的一个创新操作：在缩减特征图shape提取特征的同时，始终保留高分辨率特征图
if i < num_branches_pre:
#如果branches_cur通道数=branches_pre通道数，那么这个分支直接就可以用，不用做任何变化
#如果branches_cur通道数！=branches_pre通道数，那么就要用一个cnn网络改变通道数
#注意这个cnn是不会改变特征图的shape
#在stage1中，pre通道数是256，cur通道数为32，所以要添加这一层cnn改变通道数
#所以transition_layers第一层为
#conv2d(256,32,3,1,1)
#batchnorm2d(32)
#relu
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)
#由于branches_cur有两个分支，branches_pre只有一个分支
#所以我们必须要利用branches_pre里的分支无中生有一个新分支
#这就是常见的缩减图片shape，增加通道数提特征的操作
else:
conv3x3s = []
#这里有一个for j作用：无论stage1的分指数为多少都能顺利构建模型
#如果将stage1的分支设为3，那么需要生成2个新分支
#第一个新分支需要由branches_pre最后一个分支缩减一次shape得到
#但第二个新分支需要由branches_pre最后一个分支缩减两次shape得到，所以要做两次cnn，在第二次cnn才改变通道数
#如果stage1分支设为4也是同样的道理
#不过我们这里还是只考虑stage1分支为2的情况
for j in range(i+1-num_branches_pre):
#利用branches_pre中shape最小，通道数最多的一个分支（即最后一个分支）来形成新分支
inchannels = num_channels_pre_layer[-1]
#outchannels为64
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第二层为:
#nn.Conv2d(256, 64, 3, 2, 1, bias=False),
#nn.BatchNorm2d(64),
#nn.ReLU(inplace=True)
transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)所以self.transition1为：
(transition1): ModuleList(
(0): Sequential(
(0): Conv2d(256, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
)
(1): Sequential(
(0): Sequential(
(0): Conv2d(256, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
)
)
)而他的作用是将原来的1个分支变成两个分支：
3.4 继续加深层数，提取特征以及特征融合
我们重新回到forward
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_list里面有2个分支
#self.stage2, pre_stage_channels = self._make_stage(self.stage2_cfg, num_channels)，这里是用来做提取特征和特征融合的
#这里num_channels和上面的一样，是[32,64]
y_list = self.stage2(x_list)我们来看看self._make_stage如何实现提取特征和特征融合
文件路径：\PRTR-main\two_stage\lib\models\hrnet.py
def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES'] #1
num_branches = layer_config['NUM_BRANCHES'] #2
num_blocks = layer_config['NUM_BLOCKS'] #[4,4]
num_channels = layer_config['NUM_CHANNELS'] #[32,64]
block = blocks_dict[layer_config['BLOCK']] #BASICBLOCK
fuse_method = layer_config['FUSE_METHOD'] #SUM
modules = []
#num_modules表示一个融合块中要进行几次融合，前几次融合是将其他分支的特征融合到最高分辨率的特征图上，只输出最高分辨率特征图（multi_scale_output = False）
#只有最后一次的融合是将所有分支的特征融合到每个特征图上，输出所有尺寸特征图（multi_scale_output=True）
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
#modules第一层是 HighResolutionModule(2,BASICBLOCK,[4,4],[32,64],[32,64],SUM,reset_multi_scale_output=True)
modules.append(
HighResolutionModule(
num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output
)
)
#获取现在各个分支有多少通道
num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels我们先看看HighResolutionModule的forward函数
def forward(self, x):
#在stage1中self.num_branches为2，所以不符合if条件
#如果只有1个分支，就直接将单个分支特征图作为输入进入self.branches里设定的layers
if self.num_branches == 1:
return [self.branches[0](x[0])]
#如果有多个分支，self.branches会是一个有两个元素（这里的元素是预设的layers）的列表
#把对应的x[i]输入self.branches[i]即可
#self.branches = self._make_branches(2, BASICBLOCK, [4,4], [32,64])
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])我们再看看self._make_branches具体代码：
文件路径：\PRTR-main\two_stage\lib\models\hrnet.py
def _make_branches(self, num_branches, block, num_blocks, num_channels):
"""
并行分支的 ModuleList 结构
:param num_branches: 分支数
:param block: BASIC/BOTTLENECK
:param num_blocks: 每个分支 block 重复次数
:param num_channels: 每个分支 channel
:return:
"""
branches = []
#num_branch为2
#在stage1中branch的第一个元素为self._make_one_branch(0, BASICBLOCK, [4,4], [32,64])
#第二个元素为：self._make_one_branch(1, BASICBLOCK, [4,4], [32,64])
for i in range(num_branches):
branches.append(
self._make_one_branch(i, block, num_blocks, num_channels)
)
return nn.ModuleList(branches)self._make_one_branch代码：
文件路径：\PRTR-main\two_stage\lib\models\hrnet.py
def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
"""
一个分支的 Sequential 结构
:param branch_index: 第几个 branch
:param block: 类型
:param num_blocks: 重复次数, cfg 每个 branch 设置的次数都 = 4
:param num_channels: channel
:param stride:
:return:
"""
#这里与上面第二步的self._make_layer类似，也是一个残差结构
#这里block.expansion为1，self.num_inchannels是[32,64],num_channels[32,64]所以就不用下采样改变通道数了
# 判断是否是 stage 连接处
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第一层为：
layers.append(
block(
self.num_inchannels[branch_index],
num_channels[branch_index],
stride,
downsample
)
)
#通道数依然是[32,64]
self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
#num_blocks为[4,4]，所以有3个循环
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)返回来看HighResolutionModule的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])这一部分做的其实就是每个分支继续加深层数，提特征
在stage1中，分支1所经历的layers：
(0): Sequential(
(0): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(2): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(3): BasicBlock(
(conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)分支2：
(1): Sequential(
(0): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(1): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(2): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(3): BasicBlock(
(conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
)所以实际上实现了这个部分
3.5 特征融合
接下来看HighResolutionModule的forward函数后面部分：
x_fuse = []
#self.fuse_layers = self._make_fuse_layers()
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那么我们看看self._make_fuse_layers()代码：
文件路径：\PRTR-main\two_stage\lib\models\hrnet.py
def _make_fuse_layers(self):
"""
混合 branch 输出结果，得到 fusion 特征
:return:
fuse ModuleList(): 每个 branch 都会输出一组 生成不同大小 output 的 Sequential
[
branch1 ModuleList(), 1/4 -> [1/4, 1/8, 1/16]
branch2 ModuleList(), 1/8 -> [1/4, 1/8, 1/16]
branch3 ModuleList(), 1/16 -> [1/4, 1/8, 1/16]
]
"""
#如果只有一个分支，则不需要融合
if self.num_branches == 1:
return None
num_branches = self.num_branches #2
num_inchannels = self.num_inchannels #[32,64]
fuse_layers = []
#如果self.multi_scale_output为True，意味着只需要输出最高分辨率特征图，
#即只需要将其他尺寸特征图的特征融合入最高分辨率特征图中
#但在stage1中，self.multi_scale_output为True，所以range为2
#i表示现在要把所有分支的特征（j）融合入第i分支的特征中
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
#对于j分支进行上采样或者下采样处理，使j分支的通道数以及shape等于i分支
for j in range(num_branches):
#j > i表示j通道多于i，但shape小于i，需要上采样
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')
)
)
#j = i表示j与i为同一个分支，不需要做处理
elif j == i:
fuse_layer.append(None)
#剩余情况则是，j < i，表示j通道少于i，但shape大于i，需要下采样，利用一层或者多层conv2d进行下采样
else:
conv3x3s = []
#这个for k就是实现多层conv2d，而且只有最后一层加激活函数relu
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)重新返回HighResolutionModule的forward函数后面部分：
x_fuse = []
#现在已知self.fuse_layers里面有num_branches（上面的i）个元素fuse_layer
#接下来就把不同的x分支输入到相应的self.fuse_layers元素中分别进行上采样和下采样
#然后进行融合（相加实现融合）
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所以， y_list = self.stage2(x_list)可以实现特征融合
在stage1中layer为：
(fuse_layers): ModuleList(
(0): ModuleList(
(0): None
(1): Sequential(
(0): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): Upsample(scale_factor=2.0, mode=nearest)
)
)
(1): ModuleList(
(0): Sequential(
(0): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
(1): None
)
)
(relu): ReLU(inplace=True)
)
)返回到最初的forward函数，stage2和stage3的操作和stage1是一样的，只是像分支数这些参数有所不同
他们就是不断地增加分支、加深、融合
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_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)
y_list = self.stage2(x_list)
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])
y_list = self.stage3(x_list)
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])
y_list = self.stage4(x_list)
x = self.final_layer(y_list[0])
return x在原HRNET模型中，最后我们看self.final_layer(y_list[0])，输出关键点，至此，整个模型就结束了
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
)在PRTR模型中，最后是要生成heatmap图，通过以下的_make_head过程
def _make_head(self, pre_stage_channels):
head_block = Bottleneck
head_channels = [32, 64, 128, 256]
# Increasing the #channels on each resolution
# from C, 2C, 4C, 8C to 128, 256, 512, 1024
incre_modules = []
for i, channels in enumerate(pre_stage_channels):
incre_module = self._make_layer(head_block,
channels,
head_channels[i],
1,
stride=1)
incre_modules.append(incre_module)
incre_modules = nn.ModuleList(incre_modules)
downsamp_modules = []
for i in range(len(pre_stage_channels)-1):
in_channels = head_channels[i] * head_block.expansion
out_channels = head_channels[i+1] * head_block.expansion
downsamp_module = nn.Sequential(
nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=2,
padding=1),
nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM),
nn.ReLU(inplace=True)
)
downsamp_modules.append(downsamp_module)
downsamp_modules = nn.ModuleList(downsamp_modules)
final_layer = nn.Sequential(
nn.Conv2d(
in_channels=head_channels[3] * head_block.expansion,
out_channels=2048,
kernel_size=1,
stride=1,
padding=0
),
nn.BatchNorm2d(2048, momentum=BN_MOMENTUM),
nn.ReLU(inplace=True)
)
return incre_modules, downsamp_modules, final_layer通过以上过程生成大小为8*8的heatmap图，大体过程可以参考下图
经过多次卷积和pooling以后，得到的图像越来越小，分辨率越来越低。其中图像到 H/32∗W/32 的时候图片是最小的一层时，所产生图叫做heatmap热图，热图就是我们最重要的高维特征图。
4 数据集准备
4.1 mpii.py
文件路径：\PRTR-main\two_stage\lib\dataset\mpii.py
通过阅读源码可以知道，通过mpii.py文件中的MPIIDataset的初始化函数，将获得一个rec的数据，其中包含：MPII中所有人体，对应关键点的信息、图片路径、标准化以及缩放比例等信息。
4.1.1 _init_函数
class MPIIDataset(JointsDataset):
def __init__(self, cfg, root, image_set, is_train, transform=None):
super().__init__(cfg, root, image_set, is_train, transform)
self.num_joints = 16
self.flip_pairs = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
self.parent_ids = [1, 2, 6, 6, 3, 4, 6, 6, 7, 8, 11, 12, 7, 7, 13, 14]
self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15)
self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6)
self.db = self._get_db()
if is_train and cfg.DATASET.SELECT_DATA:
self.db = self.select_data(self.db)
logger.info('=> load {} samples'.format(len(self.db)))
MPIIDataSet类的初始化方法_init_需要如下参数：
- num_joints : MPII数据集中人体关键点标记个数
- flip_pairs : 人体水平对称关键映射
- parents_ids : 父母ids
- upper_body_ids : 定义上半身关键点
- lower_body_ids : 定义下半身关键点
- db : 读取目标检测模型
4.1.2 _get_db函数
def _get_db(self):
# create train/val split
file_name = os.path.join(
self.root, 'annot', self.image_set+'.json'
)
with open(file_name) as anno_file:
anno = json.load(anno_file)
gt_db = []
for a in anno:
image_name = a['image']
# mpii标注中的center和scale是指：
# H * W的原图像中，bbox的框原来应该是四个坐标确定，这里是用center和scale两个值来表示
# bbox的center即为center， 而bbox在mpii中默认是正方形，边长（宽） = scale * 200，这个200是官方定的
c = np.array(a['center'], dtype=np.float)
s = np.array([a['scale'], a['scale']], dtype=np.float)
# 因为mpii直接默认bbox为正方形，因此可能真正的bbox是矩形，调成正方形后可能会把人体某些部分给裁掉，所以直接把正方形扩大
if c[0] != -1:
c[1] = c[1] + 15 * s[1]
s = s * 1.25
c = c - 1
# 用到的都只有前两维
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
if self.image_set != 'test':
joints = np.array(a['joints'])
joints[:, 0:2] = joints[:, 0:2] - 1
joints_vis = np.array(a['joints_vis'])
assert len(joints) == self.num_joints, \
'joint num diff: {} vs {}'.format(len(joints),
self.num_joints)
joints_3d[:, 0:2] = joints[:, 0:2]
joints_3d_vis[:, 0] = joints_vis[:]
joints_3d_vis[:, 1] = joints_vis[:]
image_dir = 'images.zip@' if self.data_format == 'zip' else 'images'
gt_db.append(
{
'image': os.path.join(self.root, image_dir, image_name),
'center': c,
'scale': s,
'joints_3d': joints_3d,
'joints_3d_vis': joints_3d_vis,
'filename': '',
'imgnum': 0,
}
)
return gt_db
首先找到MPII数据集的分割依据文件annotaion，之后循环遍历该数据集，读取每张图片的名称、中心点位置、大小、人体关键节点位置（用三维坐标表示）、可见的人体关键节点位置并保存，形成一个字典不断加入到gt_db，循环结束返回。数据预处理到这并没有结束，因为还需要进一步处理，原因在于当计算loss的时候，我们需要的是热图(heatmap)。
4.2 JointsDataset.py
接下来，我们需要根据get_db中的信息，读取图片像素(用于训练)，同时把标签信息转化为heatmap。
文件路径：\PRTR-main\two_stage\lib\dataset\JointsDataset.py
4.2.1 init.py
class JointsDataset(Dataset):
def __init__(self, cfg, root, image_set, is_train, transform=None):
self.num_joints = 0# 人体关节的数目
self.pixel_std = 200# 像素标准化参数
self.flip_pairs = []# 水平翻转
self.parent_ids = []# 父母ID==
self.is_train = is_train# 是否进行训练
self.root = root# 训练数据根目录
self.image_set = image_set# 图片数据集名称,如‘train2017’
self.output_path = cfg.OUTPUT_DIR# 输出目录
self.data_format = cfg.DATASET.DATA_FORMAT# 数据格式如‘jpg’
self.scale_factor = cfg.DATASET.SCALE_FACTOR# 缩放因子
self.rotation_factor = cfg.DATASET.ROT_FACTOR # 旋转角度
self.flip = cfg.DATASET.FLIP# 是否进行水平翻转
self.num_joints_half_body = cfg.DATASET.NUM_JOINTS_HALF_BODY# 人体一半关键点的数目，默认为8
self.prob_half_body = cfg.DATASET.PROB_HALF_BODY# 人体一半的概率
self.color_rgb = cfg.DATASET.COLOR_RGB# 图片格式，默认为rgb
self.target_type = cfg.MODEL.TARGET_TYPE# 目标数据的类型，默认为高斯分布
self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)# 网络训练图片大小,如[192,256]
self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)# 标签热图的大小
self.sigma = cfg.MODEL.SIGMA# sigma参数，默认为2
self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT# 是否对每个关节使用不同的权重,默认为false
self.joints_weight = 1# 关节权重
self.transform = transform# 数据增强,转换等
self.db = []# 用于保存训练数据的信息,由子类提供
_init_函数的功能在于初始化JointsDataset模型，设置一些参数和参数默认值，每个参数值的作用已经注释。通过这些初始化操作，可以获得一些基本信息，如人体关节数目、图片格式、标签热图的大小、关节权重等。
4.2.2 _getitem_函数
def __getitem_(self,idx):
db_rec = copy.deepcopy(self.db[idx])
image_file = db_rec['image']
filename = db_rec['filename'] if 'fename' in db_rec else ''
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
if self.data_format == 'zip':
from utils import zipreader
data_numpy = zipreader.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
else:
data_numpy = cv2.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
if self.color_rgb:
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
if data_numpy is None:
logger.error('=> fail to read {}'.format(image_file))
raise ValueError('Fail to read {}'.format(image_file))
joints = db_rec['joints_3d']# 人体3d关键点的所有坐标
joints_vis = db_rec['joints_3d_vis']# 人体3d关键点的所有可视坐标
# 获取训练样本转化之后的center以及scale,
c = db_rec['center']
s = db_rec['scale']
# 如果训练样本中没有设置score,则加载该属性,并且设置为1
score = db_rec['score'] if 'score' in db_rec else 1
r = 0
if self.is_train:
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body
and np.random.rand() < self.prob_half_body):
c_half_body, s_half_body = self.half_body_transform(
joints, joints_vis
)
if c_half_body is not None and s_half_body is not None:
c, s = c_half_body, s_half_body
sf = self.scale_factor
rf = self.rotation_factor
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
if random.random() <= 0.6 else 0
if self.flip and random.random() <= 0.5:
data_numpy = data_numpy[:, ::-1, :]
joints, joints_vis = fliplr_joints(
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)
c[0] = data_numpy.shape[1] - c[0] - 1
trans = get_affine_transform(c, s, r, self.image_size)
input = cv2.warpAffine(
data_numpy,
trans,
(int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
if self.transform:
input = self.transform(input)
for i in range(self.num_joints):
if joints_vis[i, 0] > 0.0:
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)
target, target_weight = self.generate_target(joints, joints_vis)
target = torch.from_numpy(target)
target_weight = torch.from_numpy(target_weight)
meta = {
'image': image_file,
'filename': filename,
'imgnum': imgnum,
'joints': joints,
'joints_vis': joints_vis,
'center': c,
'scale': s,
'rotation': r,
'score': score
}
return input, target, target_weight, meta
- 首先根据idx从db获取样本信息，包括图片路径和图片序号等，如果数据格式为zip则解压，否则直接读取图像，获得像素值；再次读取db，获取人体关键点坐标、训练样本转化之后的center以及scale。
- 之后如果是进行训练，则判断可见关键点是否大于人体一半关键点，并且生成的随机数小于self.prob_half_body=0.3，如果是，则需要重新调整center和scale；再设置缩放因子和旋转因子大小，对数据进行数据增强操作，包括缩放、旋转、翻转等。
- 因为进行仿射变换，样本数据关键点发生角度旋转之后,每个像素也旋转到对应位置，所以人体的关键点也要进行仿射变换。
- 最终通过self.generate_target(joints, joints_vis)函数获得target，target_weight，shape为target[17,64,48], target_weight[17,1]。
4.2.3 cv2.warpAffine 参数详解
cv2.warpAffine 参数详解_qq878594585的博客-CSDN博客_cv2.warpaffine本文为作者原创文章，未经同意严禁转载！opencv中的仿射变换在python中的应用并未发现有细致的讲解，函数cv2.warpAffine的参数也模糊不清，今天和大家分享一下参数的功能和具体效果，如下：官方给出的参数为：cv2.warpAffine(src, M, dsize[, dst[, flags[, borderMode[, borderValue]]]]) → dst其中...https://blog.csdn.net/qq878594585/article/details/81838260
cv2.warpAffine(src, M, dsize[, dst[, flags[, borderMode[, borderValue]]]]) → dst
其中：
src - 输入图像。
M - 变换矩阵。
dsize - 输出图像的大小。
flags - 插值方法的组合（int 类型！）
borderMode - 边界像素模式（int 类型！）
borderValue - （重点！）边界填充值; 默认情况下，它为0。
上述参数中：M作为仿射变换矩阵，一般反映平移或旋转的关系，为InputArray类型的2×3的变换矩阵。
flages表示插值方式，默认为 flags=cv2.INTER_LINEAR，表示线性插值，此外还有：cv2.INTER_NEAREST（最近邻插值） cv2.INTER_AREA （区域插值） cv2.INTER_CUBIC（三次样条插值） cv2.INTER_LANCZOS4（Lanczos插值）
将之前得到的所有信息按照字典的形式，存储于meta中，并且返回input、target、 target_weight, meta。
def fliplr_joints(joints, joints_vis, width, matched_parts, pixel_align=True,
is_vis_logit=False):
"""
flip coords
"""
# Flip horizontal
joints[:, 0] = width - joints[:, 0] - int(pixel_align) # x坐标变为 w - x - 1
# Change left-right parts
for pair in matched_parts:
joints[pair[0], :], joints[pair[1], :] = \
joints[pair[1], :], joints[pair[0], :].copy()
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
if not is_vis_logit:
joints *= joints_vis
return joints, joints_vis # flip后的joint为什么还有和vis相乘我还是没搞懂？？？
4.2.4 half_body_transform函数
这个函数我觉得主要是用来数据增强的时候使用，也就是说，并不是所有的数据都是全身的关节，为了增强模型的鲁棒性，也应当适当加一些半身的图像进行训练。
def half_body_transform(self, joints, joints_vis):
# 首先获得上半身和下半身的关节id，这些关节必须都是可见的
upper_joints = []
lower_joints = []
for joint_id in range(self.num_joints):
if joints_vis[joint_id][0] > 0: # 这些关节必须都是可见的
if joint_id in self.upper_body_ids:
upper_joints.append(joints[joint_id])
else:
lower_joints.append(joints[joint_id])
# 根据概率决定是上半身还是下半身
if np.random.randn() < 0.5 and len(upper_joints) > 2:
selected_joints = upper_joints
else:
selected_joints = lower_joints \
if len(lower_joints) > 2 else upper_joints
if len(selected_joints) < 2:
return None, None
selected_joints = np.array(selected_joints, dtype=np.float32)
center = selected_joints.mean(axis=0)[:2] # 计算选出来的关节的坐标中心
# 通过右下与左上得到半身区域的宽和高来得到scale
left_top = np.amin(selected_joints, axis=0)
right_bottom = np.amax(selected_joints, axis=0)
w = right_bottom[0] - left_top[0]
h = right_bottom[1] - left_top[1]
# 保证是正方形
if w > self.aspect_ratio * h:
h = w * 1.0 / self.aspect_ratio
elif w < self.aspect_ratio * h:
w = h * self.aspect_ratio
scale = np.array(
[
w * 1.0 / self.pixel_std,
h * 1.0 / self.pixel_std
],
dtype=np.float32
)
# 适当放大，避免裁剪到人
scale = scale * 1.5
return center, scale
4.2.5 get_affine_transform函数
此函数的作用是求得仿射变换矩阵，用于下一步的关键点变换。源码的这个函数我真的看不懂，于是我把stacked hourglass network源码里进行缩放和旋转的部分代替了源码的这个函数，发现两种方法对图像的效果是一样的，所以下面我说明的是stacked hourglass network源码里的做法。这个函数我也看了特别久，原因在于之前我对仿射变换了解很少，所以建议先学习一下仿射变换以及常见的仿射变换矩阵再来看这个函数就会简单得多。
旋转后点的坐标需要通过一个旋转矩阵来确定，在网上的开源代码中，作者使用了以下矩阵的变换矩阵围绕着 (x,y) 进行任意角度的变换。
def get_affine_transform(center, scale, res, rot=0):
# Generate transformation matrix
# 首先是缩放到res尺寸
# 缩放矩阵本来应该就是[[W,0][0,H]]，但是为什么还有第三行和第三列那两个数我想了很久才想明白
h = 200 * scale[0]
t = np.zeros((3, 3))
t[0, 0] = float(res[1]) / h
t[1, 1] = float(res[0]) / h
t[0, 2] = res[1] * (-float(center[0]) / h + .5)# 把中心变到原点
t[1, 2] = res[0] * (-float(center[1]) / h + .5)# 把中心变到原点
t[2, 2] = 1
if not rot == 0:
rot = -rot # To match direction of rotation from cropping
rot_mat = np.zeros((3,3))
rot_rad = rot * np.pi / 180
sn,cs = np.sin(rot_rad), np.cos(rot_rad)
rot_mat[0,:2] = [cs, -sn]
rot_mat[1,:2] = [sn, cs]
rot_mat[2,2] = 1
# Need to rotate around center
t_mat = np.eye(3)
t_mat[0,2] = -res[1]/2
t_mat[1,2] = -res[0]/2
t_inv = t_mat.copy()
t_inv[:2,2] *= -1
t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
return t
为了更好的展示每个设置的作用，我首先把下面这两行注释掉并且把rot = 0，结果如下图所示，左边是注释前的，右边是注释后的。区别在于中心点的位置。
t[0, 2] = res[1] * (-float(center[0]) / h + .5)# 把中心变到原点
t[1, 2] = res[0] * (-float(center[1]) / h + .5)# 把中心变到原点
我再把rot = 10，结果如下图所示，左边是注释前的，右边是注释后的。区别感觉在于旋转中心点的位置。注意：rot > 0，是按逆时针旋转的。
4.2.7 generate_target函数
关键点检测主流做法还是以热图作为ground truth，通过MSE进行优化。
有关human pose estimation的问题都采用的在groundTruth坐标位置加二维高斯函数生成heatmap，从而让网络输出二维predictedHeatmap，训练后者与前者接近，最后用NMS或动态规划算法得到输出的二维坐标。
事实上，关节点检测的最终任务依然是输出预测关节点位置的坐标，然而直接让网络输出二维坐标来进行优化学习是一个极其非线性的过程，而且损失函数对权重的约束会比较弱。
那么构造heatmap实际上是构造了一个中间状态，这个heatmap有如下的一些优点：
1-可以让网络全卷积，因为输出就是2维图像，不需要全连接。
2-关节点之间（头和胸口，脖子和左右肩膀）是有很强的相关关系的。然而单独的对每一类关节点回归坐标值并不能捕捉利用这些相关关系，相反当回归heatmap时，一张输入图像对应的heatmap就存在这种相关关系，那就可以用来指导网络进行学习。简言之，头关节的回归可以帮助胸口关节，脖子关节的回归也可以帮助左右肩膀，反之亦然。
3-heatmap同样捕捉了前景（关节点）与背景的对比关系，同样可以用来指导网络进行学习。
这样，通过这条途径获得一个比较好的predictedHeatmap（易于学习，效果很好），再通过其他方法获得最终的关节点位置坐标，就是目前single person pose estimation的基本pipeline。
def generate_target(self, joints, joints_vis):
target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
target_weight[:, 0] = joints_vis[:, 0]
assert self.target_type == 'gaussian', \
'Only support gaussian map now!'
# 若生成heatmap的类型为高斯，初始化target[17,width/4(64),height/4(48)]
if self.target_type == 'gaussian':
# 生成heatmap_size大小的高斯热图
target = np.zeros((self.num_joints,
self.heatmap_size[1],
self.heatmap_size[0]),
dtype=np.float32)
tmp_size = self.sigma * 3 # 高斯半径的大小
# 为每个关键点生成热图target以及对应的热图权重target_weight
for joint_id in range(self.num_joints):
# 先计算出原图到输出热图的缩小倍数
feat_stride = self.image_size / self.heatmap_size
mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
# Check that any part of the gaussian is in-bounds 根据tmp_size参数,计算出关键点范围左上角和右下角坐标
ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
# 判断该关键点是否处于热图之外,如果处于热图之外,则把该热图对应的target_weight设置为0,然后continue
if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
or br[0] < 0 or br[1] < 0:
# If not, just return the image as is
target_weight[joint_id] = 0
continue
# # Generate gaussian 产生高斯分布的大小
size = 2 * tmp_size + 1
# x[ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12.]
x = np.arange(0, size, 1, np.float32)
# y[[ 0.][ 1.][ 2.][ 3.][ 4.][ 5.][ 6.][ 7.][ 8.][ 9.][10.][11.][12.]]
y = x[:, np.newaxis]
# x0 = y0 = 6
x0 = y0 = size // 2
# The gaussian is not normalized, we want the center value to equal 1 g形状[13,13], 该数组中间的[7,7]=1,离开该中心点越远数值越小
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))
# Usable gaussian range 判断边界,获得有效高斯分布的范围
g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
# Image range 判断边界,获得有有效的图片像素边界
img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])
# 如果该关键点对应的target_weight>0.5(即表示该关键点可见),则把关键点附近的特征点赋值成gaussian
v = target_weight[joint_id]
if v > 0.5:
target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
# 如果各个关键点训练权重不一样
if self.use_different_joints_weight:
target_weight = np.multiply(target_weight, self.joints_weight)
return target, target_weight
对于人体姿态关键点的ground truth,采用二维高斯分布，在每个关键点的ground truth位置上以1个像素为中心，生成ground truth heatmpas。首先初始化target[17,width/4,height/4]表示每个关键点的heatmap,定义heatmap的生成方式为高斯，引入关键点可见度target_weight.然后遍历每个关键点，计算其高斯heatmap,首先求出对应关键点在heatmap中的坐标，这里定义feat_stride=4,即heatmap比原图的尺寸缩小4倍。然后检查对于此位置执行高斯操作时，是否会超过界限，若超出则将关键点的可见度置为0，跳出此次循环。若在范围内，则定义高斯size,执行高斯函数，将中心值设置为1。接着，分别计算出高斯和heatmap的有效范围，若此关键点的可见度大于0.5，则将对应区域的高斯值赋给heatmap。所以最后计算出了所有关键点的heatmap,若设置不同关键点加权，则执行权重与可见度相乘。
该函数的功能主要在于产生热图，并且制作热图的方式必须为gaussion。会为每个关键点生成热图target以及对应的热图权重target_weight，在生成期间还需判断该关键点是否处于热图之外,如果处于热图之外,则把该热图对应的target_weight设置为0,然后continue。最终生成高斯分布的热图表示，返回target和target_weight。
4.3 数据增强
transform是对bbox进行的，不是对原图像，因此要注意center的位置，要进行相应的平移把bbox移到想要进行的transform对应的初始坐标处。
总共分为两步：
- 缩放与平移：将原图坐标缩放到输出尺寸，然后将左上角变到原点
- 旋转：首先将轴心(x,y)移到原点，然后做旋转平移变换，最后再将图像的左上角转换为原点
4.3.1 缩放与平移
res的shape是(H, W)
变换过程可以用下式表示：
读入数据后，需要先把大小不一的标注图片统一转换成 256 x 256。
对于 LSP 测试集，作者使用的是图像的中心作为身体的位置，并直接以图像大小来衡量身体大小。数据集里的原图片是大小不一的（原图尺寸存在 bbox 里），一般采取 crop 的方法有好几种，比如直接进行 crop，然后放大，这样做很明显会有丢失关节点的可能性。也可以先把图片放在中间，然后将图片缩放到目标尺寸范围内原尺寸的可缩放的大小，然后四条边还需要填充的距离，最后 resize 到应有大小。
这里采用的是先扩展边缘，然后放大图片，再进行 crop，这样做能够保证图片中心处理后依然在中心位置，且没有关节因为 crop 而丢失。注意在处理图片的同时需要对标注也进行处理。
4.3.2 旋转
为什么要把中心移来移去？缩放变换的矩阵中心随意，只要把对应的W和H确定好就行，但是旋转就有中心一说了。我们想要缩放后的框按中心点旋转，旋转矩阵常见的起始点是原点，也即
[cos, -sin]
[sin, cos]
所以把缩放后的框移到对应的位置，就可以利用这个矩阵进行旋转了，当然也可以不移动，但是对应的旋转矩阵就要进行相应的变换，我只是解释一下源码的做法。
这里需要注意一下，图像的坐标轴和我们平时画的不一样（y轴的方向不一样）所以上面的矩阵在我们正常的坐标系里是逆时针，但在图像的坐标轴里，是我们人眼认知的顺时针。所以解释了源码里的这一句：
rot = -rot # To match direction of rotation from cropping
源码使用的旋转矩阵是正常情况下的逆时针旋转矩阵，那么会使图像顺时针转动，但是源码想要图像逆时针旋转r°，所以就把rot = -rot就变成了逆时针。
此步可以总结为以下变换矩阵。矩阵由来：首先将轴心(x,y)移到原点，然后做旋转缩放变换，最后再将图像的左上角转换为原点
def affine_transform(pt, t):
#把对应的gt也进行相应的transform
# 2,3 * 3, --->2,
new_pt = np.array([pt[0], pt[1], 1.]).T
new_pt = np.dot(t, new_pt)
return new_pt[:2]
除此之外，以下数据增强的方法也很常见：
1. 从颜色上考虑，还可以做图像亮度、饱和度、对比度变化、PCA Jittering（按照 RGB 三个颜色通道计算均值和标准差后在整个训练集上计算协方差矩阵，进行特征分解，得到特征向量和特征值）；
2. 从图像空间性质上考虑，还可以使用随机裁剪、平移；
3. 从噪声角度，高斯噪声、椒盐噪声、模糊处理；
4. 从类别分布的角度，可以采用 label shuffle、Supervised Data Augmentation（海康威视 ILSVRC 2016 的 report）。
4.3.3 典型的仿射变换
5 HRNET模型设计
5.1 基本模块
如下的左图对应于resnet-18/34使用的基本块，右图是50/101/152所使用的，由于他们都比较深，所以有图相比于左图使用了1x1卷积来降维。
基本模块主要是BasicBlock、Bottleneck，现在进行逐个分析：
5.1.1 BasicBlock ： 搭建上图左边的模块。
- 每个卷积块后面连接BN层进行归一化；
- 残差连接前的3x3卷积之后只接入BN，不使用ReLU，避免加和之后的特征皆为正，保持特征的多样；
- 跳层连接：两种情况，当模块输入和残差支路（3x3->3x3）的通道数一致时，直接相加；当两者通道不一致时（一般发生在分辨率降低之后，同分辨率一般通道数一致），需要对模块输入特征使用1x1卷积进行升/降维（步长为2，上面说了分辨率会降低），之后同样接BN，不用ReLU。
# 提前写好一个类，在HighResolutionModule中使用
# x: [batch_size, 256, 8, 8]
# out: [batch_size, 256, 8, 8]
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
5.1.2 Bottleneck ：搭建上图右边的模块。
- 使用1x1卷积先降维，再使用3x3卷积进行特征提取，最后再使用1x1卷积把维度升回去；
- 每个卷积块后面连接BN层进行归一化
- 残差连接前的1x1卷积之后只接入BN，不使用ReLU，避免加和之后的特征皆为正，保持特征的多样性。
- 跳层连接：两种情况，当模块输入和残差支路（1x1->3x3->1x1）的通道数一致时，直接相加；当两者通道不一致时（一般发生在分辨率降低之后，同分辨率一般通道数一致），需要对模块输入特征使用1x1卷积进行升/降维（步长为2，上面说了分辨率会降低），之后同样接BN，不用ReLU。
# 在layer1中使用4个Bottleneck，验证论文中以高分辨率子网为第一阶段，维持高分辨率表示
# x: [batch_size, 256, 64, 64]
# output : [batch_size, 256, 64, 64]
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
5.2 高分辨率模块-HighResolutionModule
根据ResNet的设计规则，将深度分布到每个阶段，将通道数分布到每个分辨率，实例化了关键点热图估计网络。
HRNet，包含四个阶段，四个并行的子网，其分辨率逐渐降低到一半，相应的宽度(通道的数量)增加到原来的两倍。第一阶段包含4残差单位,每个单元和ResNet-50相同,由一个宽度64的瓶颈组成，紧接着一个3×3卷积将特征图的宽度减少到c，第二,第三,第四阶段分别包含1、4、3个交换块（多分辨率模块），四个分辨率卷积层的宽度(通道数)分别是C，2C，4C和8C。一个交换块（多分辨率模块）包含4个残差单元，其中每个单元在每个分辨率中包含两个3×3的卷积，一个交换单元跨分辨率。综上所述，共有8个交换单元，即，进行了8次多尺度融合。
在实验中，研究了一个小网络和一个大网络:HRNet-W32和HRNet-W48，其中32和48分别代表高分辨率子网在最后三个阶段的宽度(C)。其他三个并行子网的宽度分别为HRNet-W32的64、128、256和HRNet-W48的96、192、384。
当仅包含一个分支时，生成该分支，没有融合模块，直接返回；当包含不仅一个分支时，先将对应分支的输入特征输入到对应分支，得到对应分支的输出特征；紧接着执行融合模块。
class HighResolutionModule(nn.Module):
def __init__(self,
num_branches,
block, num_blocks,
num_inchannels, num_channels,
fuse_method, # sum / cat
multi_scale_output=True):
"""
1.构建 branch 并行 多 scale 特征提取
2.在 module 末端将 多 scale 特征通过 upsample/downsample 方式，并用 sum 进行 fuse
注意：这里的 sum fuse 是值从 多个 branch(j) 到 branch_i 的聚合结果;
整个 module 的输出结果依然是 并行的 num_branch 个结果
:param num_branches: stage 并行高度
:param block: BASIC/BOTTLENECK
:param num_blocks: 指定每个 block 重复次数
:param num_inchannels: 由 NUM_CHANNELS 和 block.expansion 相乘得到
:param num_channels:
:param fuse_method: sum / cat
:param multi_scale_output:
"""5.2.1 _check_branches函数
判断num_branches (int) 和 num_blocks, num_inchannels, num_channels (list) 三者的长度是否一致，否则报错；
# 判断三个参数长度是否一致，否则报错
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)
5.2.2 _make_one_branch函数
搭建一个分析，单个分支内部分辨率相等，一个分支由num_branches[branch_index]个block组成，block可以是两种ResNet模块中的一种；
- 首先判断是否降维或者输入输出的通道(num_inchannels[branch_index]和 num_channels[branch_index] * block.expansion(通道扩张率))是否一致，不一致使用1z1卷积进行维度升/降，后接BN，不使用ReLU；
- 顺序搭建num_blocks[branch_index]个block，第一个block需要考虑是否降维的情况，所以单独拿出来，后面1 到 num_blocks[branch_index]个block完全一致，使用循环搭建就行。此时注意在执行完第一个block后将num_inchannels[branch_index]重新赋值为num_channels[branch_index] * block.expansion。
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)
5.2.3 _make_branches函数
循环调用_make_one_branch函数创建多个分支；
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)
5.2.4 _make_fuse_layers函数
HighResolutionModule 末尾的特征融合层
以下图红框即 stage3 中 蓝色 branch 输出结果为例，其输出结果要转换成 4 种尺度的特征，用于每个 branch 末尾的特征融合
- 1/8 ↗ 1/4，不同层，channel 不同，size 不同 通道转换 + 上采样 (在 forward 函数中由双线性插值完成)
- 1/8 → 1/8，相同层，channel 一致，size 一致 None，直接使用 feature
- 1/8 ↘ 1/16，不同层，channel 不同，size 不同 通道转换 + 下采样 (通过串联的 stride=2 的 3x3 conv 完成)
- 1/8 ↘ 1/32，同上
(1) 如果分支数等于1，返回none，说明此时不需要使用融合模块
if self.num_branches == 1:
return None
(2) 双层循环
for i in range(num_branches if self.multi_scale_output else 1)
该语句的作用在于，如果需要产生多分辨率的结果，就双层循环num_branches次，如果只需要产生最高分辨率的表示，就将i确定为0。
- 如果j > i，此时的目标是将所有分支上采样到和i分支相同的分辨率并融合，也就是说j所代表的分支分辨率比i分支低，2**(j-i)表示j分支上采样这么多倍才能和i分支分辨率相同。先使用1x1卷积将j分支的通道数变得和i分支一致，进而跟着BN，然后依据上采样因子将j分支分辨率上采样到和i分支分辨率相同，此处使用最近邻插值；
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')
)
)
- 如果
j = i，也就是说自身与自身之间不需要融合，nothing to do；
elif j == i:
fuse_layer.append(None)
- 如果j < i，转换角色，此时最终目标是将所有分支采样到和i分支相同的分辨率并融合，注意，此时j所代表的分支分辨率比i分支高，正好和(2.1)相反。此时再次内嵌了一个循环，这层循环的作用是当i-j > 1时，也就是说两个分支的分辨率差了不止二倍，此时还是两倍两倍往上采样，例如i-j = 2时，j分支的分辨率比i分支大4倍，就需要上采样两次，循环次数就是2；
i. 当k == i - j - 1时，举个例子，i = 2,j = 1, 此时仅循环一次，并采用当前模块，此时直接将j分支使用3x3的步长为2的卷积下采样(不使用bias)，后接BN，不使用ReLU；
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)
)
)
ii. 当k != i - j - 1时，举个例子，i = 3,j = 1, 此时循环两次，先采用当前模块，将j分支使用3x3的步长为2的卷积下采样(不使用bias)两倍，后接BN和ReLU，紧跟着再使用(2.3.1)中的模块，这是为了保证最后一次二倍下采样的卷积操作不使用ReLU，猜测也是为了保证融合后特征的多样性；
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)
)
)
5.3 整合模块-PoseHighResolutionNet
5.3.1 stage1
进行一系列的卷积操作，获得最初的特征图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)
5.3.2 stage2
- 首先根据原先设定，获得相关配置信息。对于第二阶段，num_channels=[32,64]，num_channels表示输出通道，最后的64是新建平行分支N2的输出通道数；block为Bottleneck，在论文中提到，第一个stage到第二个stage变换时，使用Bottleneck.
- 之后会生成新的平行N2分支网络，即N11 --> N21,N22这个过程，同时如果输入输出通道不一致时。会对输入的特征图x进行通道变换.
- 最后对平行子网进行加工，让其输出的y，可以当做下一个stage的输入x，这里的pre_stage_channels为当前阶段的输出通道数，也就是一个stage的输入通道数，同时平行子网信息交换模块，也包含其中
self.stage2_cfg = extra['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition1 = self._make_transition_layer([256], num_channels)
self.stage2, pre_stage_channels = self._make_stage(
self.stage2_cfg, num_channels)
5.3.3 stage 3
- 首先根据原先设定，获取stage3的相关配置信息。对于第三阶段，num_channels=[32,64,128],num_channels表示输出通道,最后的128是新建平行分支N3的输出通道数；这里的block为BasicBlock,在论文中有提到,除了第一个stage到第二个stage变换时使用Bottleneck,其余的都是使用BasicBlock
- 之后会生成新的平行分支N3网络,即N22–>N32,N33这个过程时，如果输入输出通道不一致时。会对输入的特征图x进行通道变换.
- 最后对平行子网进行加工，让其输出的y，可以当做下一个stage的输入x，这里的pre_stage_channels为当前阶段的输出通道数，也就是一个stage的输入通道数，同时平行子网信息交换模块，也包含其中
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)
5.3.4 stage 4
- 首先根据原先设定，获取stage4的相关配置信息。对于第四阶段，num_channels=[32,64,128,256],num_channels表示输出通道,最后的256是新建平行分支N3的输出通道数；这里的block为BasicBlock
- 之后会生成新的平行分支N3网络,即N33–>N43,N44这个过程时，如果输入输出通道不一致时。会对输入的特征图x进行通道变换.
- 最后对平行子网进行加工，让其输出的y，可以当做下一个stage的输入x，这里的pre_stage_channels为当前阶段的输出通道数，也就是一个stage的输入通道数，同时平行子网信息交换模块，也包含其中
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)
5.3.5 整合预测
对最终的特征图混合之后，进行一次卷积，预测人体关键点的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
)
self.pretrained_layers = extra['PRETRAINED_LAYERS']
5.3.6 _make_transition_layer
该函数的作用在于生成下一阶段同等分辨率和一般分辨率的分支。首先会进行循环遍历，对每个分支进行处理
不是最后一个分支：如果当前一层的输入通道和输出通道不相等，则通过卷积对通道数进行变换；如果当前层的输入=输出通道数，则维持原样；
是最后一个分支：新建一个分支，并且这个分支分辨率会减少一半
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)
"""
:param num_channels_pre_layer: pre_stage output channels list
:param num_channels_cur_layer: cur_stage output channels list
cur 总比 pre 多一个 output_channel 对应增加的 1/2 下采样
stage2 stage3 stage4
pre: [256] [48,96] [48,96,192]
cur: [48,96] [48,96,192] [48,96,192,384]
每个 stage channels 数量也对应了 stage2/3/4 使用 BASIC block; expansion=1
:return:
transition_layers:
1.完成 pre_layer 到 cur_layer input channels 数量对应
2.完成 feature map 尺寸对应
"""
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]),
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)
以下为 hrnet_w48 的 transition 具体结构
# stage 1-2
(transition1): ModuleList(
# input channels，从 1/4 到 1/4，完成通道数量转换
(0): Sequential(
(0): Conv2d(256, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(1): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
# input channels + downsample，从 1/4 到 1/8，不仅通道数量，而且使用 stride=2 进行下采样
(1): Sequential(
(0): Sequential(
(0): Conv2d(256, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
)
# stage 2-3
(transition2): ModuleList(
(0): None # 因为 同层对应的连接处的 feature map channels 和 size 一致，所以不需要转换
(1): None
# downsample，stage2 末尾，从 1/8 到 1/16，需要使用 stride=2 下采样
(2): Sequential(
(0): Sequential(
(0): Conv2d(96, 192, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
)
# stage 3-4
(transition3): ModuleList(
(0): None
(1): None
(2): None
# downsample，同 stage2 用法一样，因为前3个branch对应的 feature map 可以直接连接，所以只要对末尾完成 1/16 到 1/32 下采样
(3): Sequential(
(0): Sequential(
(0): Conv2d(192, 384, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
)5.3.7 _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']
# 当stage=2,3,4时,分别为：BasicBlock,BasicBlock,
block = blocks_dict[layer_config['BLOCK']]
# 当stage=2,3,4时,都为：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
# multi_scale_output 只被用再最后一个HighResolutionModule
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
5.4 forward
5.4.1 第一阶段
经过一系列的卷积, 获得初步特征图,总体过程为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)
5.3.2 第二阶段
其中包含了创建分支的过程,即 N11–>N21,N22
总体过程为:
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)
y_list = self.stage2(x_list)
5.3.3 第三阶段
其中包含了创建分支的过程,即 N22–>N32,N33
总体过程为：
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])
y_list = self.stage3(x_list)
5.3.4 第四阶段
其中包含了创建分支的过程,即 N33–>N43,N44
总体过程为：
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]
5.3.5 预测阶段
y[b, 32, 64, 48] --> x[b, 16, 64, 48]
x = self.final_layer(y_list[0])
6 transformer模型设计
7 源码分析（训练阶段）
解析参数->构建网络模型->加载训练测试数据集迭代器->迭代训练->模型评估保存
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='')
# log 输出tensorboard的目录
parser.add_argument('--logDir',
help='log directory',
type=str,
default='')
# data 训练数据的目录
parser.add_argument('--dataDir',
help='data directory',
type=str,
default='')
# premodel 预训练模型的目录
parser.add_argument('--prevModelDir',
help='prev Model directory',
type=str,
default='')
args = parser.parse_args()
return args
def main():
args = parse_args() # 对输入的参数进行解析
update_config(cfg, args) # 根据输入参数对cfg进行更新
# 创建logger，用于记录训练过程的打印信息
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
#根据配置文件构建网络
print('models.'+cfg.MODEL.NAME+'.get_pose_net')
model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
cfg, is_train=True
)
# copy model file 拷贝lib/models/pose_hrnet.py文件到输出目录之中
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))
# 让模型支持多GPU训练
model = torch.nn.DataParallel(model, device_ids=[0,]).cuda()
# 用于计算loss
# 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=1,
shuffle=cfg.TRAIN.SHUFFLE,
num_workers=0,
pin_memory=cfg.PIN_MEMORY
)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=1,
shuffle=False,
num_workers=0,
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()
8 源码分析（测试阶段）
人体姿态估计推导inference的过程如下。首先读入验证集数据，初始化预测关键点结果的矩阵和人体姿态边界框。在这个过程中将batch归一化和dropout都禁用。控制数据不计算梯度和反向传播。将所有数据导入设备当中。调用model分别对原始图片和反转后的图片执行正向传播过程，将所得到heatmap结果进行融合。然后计算output与target之间的loss,和进行评估得到准确率值。
对于在HRNet model中具体的执行步骤如下图。输入图片首先通过两组conv+bn+relu操作，然后建立新的分支得到两个子网，在stage2中有进行1次信息交换；再建立新的分支得到三个子网，在stage3中有4次信息交换；建立最后一个分支，得到4个子网，在stage4中有3次信息交换，最后通过卷积层输出。在建立分支过程中，特征分辨率逐渐减半。
9 亲测结果
Results on MPII val：绿色背景部分表示已亲测验证，详细结果见下标题截图
| Arch | Input Size | Head | Shoulder | Elbow | Wrist | Hip | Knee | Ankle | Mean | [email protected] |
|---|---|---|---|---|---|---|---|---|---|---|
| prtr_res50 | 256x256 | 94.577 | 93.088 | 83.109 | 74.079 | 84.733 | 74.732 | 69.367 | 82.865 | 22.602 |
| prtr_res50 | 384x384 | 96.010 | 94.124 | 86.586 | 79.940 | 86.464 | 81.604 | 74.563 | 86.310 | 28.704 |
| prtr_res50 | 512x512 | 96.555 | 95.007 | 88.597 | 83.383 | 88.731 | 84.002 | 79.121 | 88.493 | 34.036 |
| prtr_res101 | 256x256 | 94.816 | 93.461 | 84.728 | 76.853 | 87.000 | 79.730 | 72.768 | 84.975 | 25.517 |
| prtr_res101 | 384x384 | 96.282 | 94.990 | 88.307 | 82.353 | 88.125 | 83.639 | 77.445 | 87.916 | 31.642 |
| prtr_res101 | 512x512 | 96.828 | 95.839 | 90.234 | 84.633 | 89.302 | 85.049 | 80.043 | 89.409 | 37.198 |
| prtr_res152 | 256x256 | 96.146 | 94.480 | 86.108 | 78.515 | 87.658 | 81.826 | 74.634 | 86.313 | 26.885 |
| prtr_res152 | 384x384 | 96.419 | 94.871 | 88.444 | 82.627 | 88.575 | 84.143 | 78.365 | 88.215 | 32.160 |
| prtr_hrnet_w32 | 256x256 | 96.794 | 95.584 | 89.603 | 83.143 | 88.731 | 83.739 | 78.012 | 88.584 | 33.206 |
| prtr_hrnet_w32 | 384x384 | 97.340 | 96.009 | 90.557 | 84.479 | 89.700 | 85.533 | 78.956 | 89.526 | 35.410 |
9.1 prtr_hrnet_w32 256x256
9.2 prtr_hrnet_w32 384x384
9.3 prtr_res152 384x384
9.3 prtr_res152 256x256
10 参考链接
HRNet_qq_39862223的博客-CSDN博客论文链接：https://arxiv.org/abs/1902.09212代码链接：https://github.com/leoxiaobin/deep-high-resolution-net.pytorch论文源码分析：1 源码准备在指定文件夹下，输入命令：git clone https://github.com/leoxiaobin/deep-high-resolution-net.pytorch.git下载完成后，得到HRNet源码2 源码结构下表列出HRNet中比较重要的文件：.https://blog.csdn.net/qq_39862223/article/details/109360652?spm=1001.2014.3001.5506hrnet模型源代码详解_绵绵是一只鹅呀的博客-CSDN博客_hrnet代码（仅为个人学习笔记，如果有错误欢迎提出）请对照源代码和本帖子观看：https://github.com/leoxiaobin/deep-high-resolution-net.pytorch建议先看看论文大概了解hrnet特点再看我们先看看代码里用来搭建模型的方法：def get_pose_net(cfg, is_train, **kwargs): model = PoseHighResolutionNet(cfg, **kwargs) if is_train and chttps://blog.csdn.net/qq_36382582/article/details/119541890?spm=1001.2014.3001.5506https://segmentfault.com/a/1190000019167646
https://segmentfault.com/a/1190000019167646[论文阅读]HRNetV1,HRNetV2,HRNetV2p_gefeng1209的博客-CSDN博客_hrnet图像分类1.Deep High-Resolution Representation Learning for Human Pose Estimation(HRNetV1)2.High-Resolution Representations for Labeling Pixels and Regions(HRNetV2,HRNetV2p)1.Introduction人体姿势估计（又称关键点检测）旨...https://blog.csdn.net/gefeng1209/article/details/93142916
HRNet 源代码结构详解 - 简书Github: https://github.com/HRNet/HRNet-Semantic-Segmentation[https://github.com/HRNet/H...https://www.jianshu.com/p/7e55b80614a7
关键点检测一：HRNet数据预处理（MPII）_Mr.Pru-CSDN博客_hrnet关键点检测关键点检测一：HRNet数据预处理（MPII）前言HRNet源码数据预处理目录位置源码分析JointsDataset类`__init__``_get_db``half_body_transform``__getitem__``get_affine_transform``generate_target`前言最近在做考场行为分析的一个项目，其中我负责的是使用关键点检测算法来进行考生异常行为检测。之前只接触过分类算法，写与看的代码也只限于分类任务。而检测任务工程量太大，因此在看官方源码时非常的吃力，因此希望写 https://blog.csdn.net/qq_43312130/article/details/122034420
11 仿射变换
如何通俗地讲解「仿射变换」这个概念？ - 知乎简单来说，“仿射变换”就是：“线性变换”+“平移”。先看什么是线性变换？1 线性变换线性变换从几何直…
https://www.zhihu.com/question/20666664