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)

  1. 通过指定 cfg_path 选择要使用的模型的结构（yaml 存储）
  2. 通过指定 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_configcfg = update_config(cfg_path)model = HighResolutionNet(cfg, **kwargs)if kwargs.get('use_pretrain', False):model.load_pretrain(cfg.MODEL.PRETRAINED)return model

yaml 文件中，关于模型结构的关键部分，以 hrnet_w48 为例

MODEL:NAME: seg_hrnetALIGN_CORNERS: TruePRETRAINED: 'pretrained_models/hrnetv2_w48_imagenet_pretrained.pth'  # 指定 pretrain 模型路径EXTRA:  # EXTRA 具体定义了模型的结果，包括 4 个 STAGE，各自的参数FINAL_CONV_KERNEL: 1STAGE1:NUM_MODULES: 1NUM_RANCHES: 1BLOCK: BOTTLENECKNUM_BLOCKS:- 4NUM_CHANNELS:- 64FUSE_METHOD: SUMSTAGE2:NUM_MODULES: 1    # HighResolutionModule 重复次数NUM_BRANCHES: 2   # 分支数BLOCK: BASICNUM_BLOCKS:- 4- 4NUM_CHANNELS:- 48- 96FUSE_METHOD: SUMSTAGE3:NUM_MODULES: 4NUM_BRANCHES: 3BLOCK: BASICNUM_BLOCKS:- 4- 4- 4NUM_CHANNELS:- 48- 96- 192FUSE_METHOD: SUMSTAGE4:NUM_MODULES: 3NUM_BRANCHES: 4BLOCK: BASICNUM_BLOCKS:- 4- 4- 4- 4NUM_CHANNELS:- 48- 96- 192- 384FUSE_METHOD: SUM

2 源码结构

下表列出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))#经过第一层后，当前特征图通道数为256self.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 = 4def __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 = downsampleself.stride = stridedef forward(self, x):residual = xout = 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: 1NUM_BRANCHES: 2BLOCK: BASICNUM_BLOCKS:- 4- 4NUM_CHANNELS:- 32- 64FUSE_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_layerself.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、1for 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)#reluif 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为64outchannels = num_channels_cur_layer[i] if j == i-num_branches_pre else inchannelsconv3x3s.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']    #1num_branches = layer_config['NUM_BRANCHES']    #2num_blocks = layer_config['NUM_BLOCKS']    #[4,4]num_channels = layer_config['NUM_CHANNELS']    #[32,64]block = blocks_dict[layer_config['BLOCK']]   #BASICBLOCKfuse_method = layer_config['FUSE_METHOD'] #SUMmodules = []#num_modules表示一个融合块中要进行几次融合，前几次融合是将其他分支的特征融合到最高分辨率的特征图上，只输出最高分辨率特征图（multi_scale_output = False）#只有最后一次的融合是将所有分支的特征融合到每个特征图上，输出所有尺寸特征图（multi_scale_output=True）for i in range(num_modules):# multi_scale_output is only used last moduleif not multi_scale_output and i == num_modules - 1:reset_multi_scale_output = Falseelse: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里设定的layersif 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 = Noneif 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 Nonenum_branches = self.num_branches  #2num_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，而且只有最后一层加激活函数relufor 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 = Bottleneckhead_channels = [32, 64, 128, 256]# Increasing the #channels on each resolution# from C, 2C, 4C, 8C to 128, 256, 512, 1024incre_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.expansionout_channels = head_channels[i+1] * head_block.expansiondownsamp_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 = 16self.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 splitfile_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.25c = 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] - 1joints_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# 人体一半关键点的数目，默认为8self.prob_half_body = cfg.DATASET.PROB_HALF_BODY# 人体一半的概率self.color_rgb = cfg.DATASET.COLOR_RGB# 图片格式，默认为rgbself.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参数，默认为2self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT# 是否对每个关节使用不同的权重,默认为falseself.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 zipreaderdata_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,则加载该属性,并且设置为1score = db_rec['score'] if 'score' in db_rec else 1r = 0if self.is_train:if (np.sum(joints_vis[:, 0]) > self.num_joints_half_bodyand 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_bodysf = self.scale_factorrf = self.rotation_factors = 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 0if 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] - 1trans = 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 horizontaljoints[:, 0] = width - joints[:, 0] - int(pixel_align)  # x坐标变为 w - x - 1# Change left-right partsfor 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_visreturn 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_jointselse:selected_joints = lower_joints \if len(lower_joints) > 2 else upper_jointsif len(selected_joints) < 2:return None, Noneselected_joints = np.array(selected_joints, dtype=np.float32)center = selected_joints.mean(axis=0)[:2] # 计算选出来的关节的坐标中心# 通过右下与左上得到半身区域的宽和高来得到scaleleft_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_ratioelif w < self.aspect_ratio * h:w = h * self.aspect_ratioscale = np.array([w * 1.0 / self.pixel_std,h * 1.0 / self.pixel_std],dtype=np.float32)# 适当放大，避免裁剪到人scale = scale * 1.5return 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]) / ht[1, 1] = float(res[0]) / ht[0, 2] = res[1] * (-float(center[0]) / h + .5)# 把中心变到原点t[1, 2] = res[0] * (-float(center[1]) / h + .5)# 把中心变到原点t[2, 2] = 1if not rot == 0:rot = -rot # To match direction of rotation from croppingrot_mat = np.zeros((3,3))rot_rad = rot * np.pi / 180sn,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 centert_mat = np.eye(3)t_mat[0,2] = -res[1]/2t_mat[1,2] = -res[0]/2t_inv = t_mat.copy()t_inv[:2,2] *= -1t = 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_weightfor joint_id in range(self.num_joints):# 先计算出原图到输出热图的缩小倍数feat_stride = self.image_size / self.heatmap_sizemu_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,然后continueif 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 istarget_weight[joint_id] = 0continue# # 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 = 6x0 = 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(即表示该关键点可见),则把关键点附近的特征点赋值成gaussianv = 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对应的初始坐标处。

总共分为两步：

  1. 缩放与平移：将原图坐标缩放到输出尺寸，然后将左上角变到原点
  2. 旋转：首先将轴心(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.]).Tnew_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 = 1def __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 = downsampleself.stride = stridedef forward(self, x):residual = xout = 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 += residualout = 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 = 4def __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 = downsampleself.stride = stridedef forward(self, x):residual = xout = 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 += residualout = 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 / catmulti_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 = Noneif 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.expansionfor 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 末尾的特征融合