【语义分割】Stacked Hourglass Networks 以及 PyTorch 实现
Stacked Hourglass Networks（级联漏斗网络）
姿态估计（Pose Estimation）是 CV 领域一个非常重要的方向，而级联漏斗网络的提出就是为了提升姿态估计的效果，但是其中的经典思想可以扩展到其他方向，比如目标识别方向，代表网络是 CornerNet（预测目标的左上角和右下角点，再进行组合画框）。
CNN 之所以有效，是因为它能自动提取出对分类、检测和识别等任务有帮助的特征，并且随着网络层数的增加，所提取的特征逐渐变得抽象。以人脸识别为例，低层卷积网络能够提取出一些简单的特征，如轮廓；中间卷积网络能够提取出抽象一些的特征，如眼睛鼻子；较高层的卷积网络则能提取出更加抽象的特征，比如完整的人脸。这些将有助于我们理解级联漏斗模型（Stacked Hourglass Model，简称SHM）为什么有效。
做姿态估计，需要预测身体不同的关节点，手臂这种线条简单的结构，可能在中间卷积网络更容易被识别；而面部这种线条复杂的结构，可能在高层卷积网络才更容易被识别。因此，如果我们只使用最后一层的 feature map，就会造成一些信息的丢失。SHN 的主要贡献——利用多尺度特征来识别姿态。
Single Hourglass Network
上图是单个漏斗网络的结构。该结构与全卷积网络和其它设计（以多尺度方式处理空间信息，并进行密集预测）紧密相连。然而漏斗网络与其它设计有什么不同呢？由图可以看出，其自底向上（从高分辨率到低分辨率）处理和自顶向下（从低分辨率到高分辨率）处理之间的容量分布（这里实在不知道怎么翻译。。。）更加对称。另外还有一点需要注意，在自顶向下处理过程中，使用的不是 unpooing（一种常见的上采样操作）或者 deconv layers（可称为去卷积层），而是采用nearest neighbor upsampling（最近邻上采样）和 skip connections。这些操作需要在源码中理解。
Stacked Hourglass Networks
上图是单个漏斗网络后面的一些设计以及两个漏斗网络的连接细节。
块1 是上面介绍的单个沙漏网络，在它后面是一个 1$\times\(1 的全卷积网络，即块2；块2 后面分离出上下两个分支（块3 和块4）：上分支（块3）依然是一个 1\)\times$1 的全卷积网络，下分支（块4）为 Heat map（下面重点介绍）。块5 是对块4 进行 channal 上的扩增，以方便块3、块5 和 上个漏斗网络的输出进行合并，一起作为当前漏斗网络的输出，同时是下一个漏斗网络的输入。
这里对 Heat map 进行解释：大部分姿态检测的最后一步是对 feature map 上的每个像素做概率预测，计算该像素是某个关节点的概率，而这里的 feature map 就是上面输出的 Heat map。使用它与真值进行误差计算。应用中，如果多个 Hourglass Module 组合在一起进行梯度下降，输出层的误差经过多层反向传播会大幅减小，也就是发生了梯度消失。因此，在整个网络中每个Hourglass Module 后面都会输出 Heat map 来计算损失。这种方法称为 中间监督（Intermediate Supervision），可以保证底层参数正常更新。
之所以使用多个 Stack Hourglass，是为了重复自下而上和自上而下的推理机制，允许重新评估整个图像的初始估计和特征，实现这一过程的核心就是预测中间的 Heat map，并让中间 Heat map 参与 loss 计算。
PyTorch 实现 Model
-
首先定义残差网络的基本模块：
import torch.nn as nn class HgResBlock(nn.Module): def __init__(self, inplanes, outplanes, stride=1): super(HgResBlock, self).__init__() self.inplanes = inplanes self.outplanes = outplanes midplanes = outplanes // 2 self.bn_1 = nn.BatchNorm2d(inplanes) self.conv_1 = nn.Conv2d(inplanes, midplanes, kernel_size=1, stride=stride) self.bn_2 = nn.BatchNorm2d(midplanes) self.conv_2 = nn.Conv2d(midplanes, midplanes, kernel_size=3, stride=1, padding=1) self.bn_3 = nn.BatchNorm2d(midplanes) self.conv_3 = nn.Conv2d(midplanes, outplanes, kernel_size=1, stride=1) self.relu = nn.ReLU(inplanes=True) if inplanes != outplanes: self.conv_skip = nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=1) # Bottle neck def forward(self, x): residual = x out = self.bn_1(x) out = self.conv_1(out) out = self.relu(out) out = self.bn_2(out) out = self.conv_2(out) out = self.relu(out) out = self.bn_3(out) out = self.conv_3(out) out = self.relu(out) if self.inplanes != self.outplanes: residual = self.conv_skip(residual) out += residual return out -
定义单个的 Hourglass Module（注意这里用到了递归）：
import torch.nn as nn class Hourglass(nn.Module): def __init__(self, depth, nFeat, nModules, resBlocks): super(Hourglass, self).__init__() self.depth = depth self.nFeat = nFeat self.nModules = nModules self.resBlocks = resBlocks self.hg = self._make_hourglass() self.downsample = nn.MaxPool2d(kernel_size=2, stride=2) self.upsample = nn.Upsample(scale_factor=2, mode='nearest') def _make_residual(self, n): return nn.Sequential(*[self.resBlocks(self.nFeat, self.nFeat) for _ in range(n)]) def _make_hourglass(self): hg = [] for i in range(self.depth): res = [self._make_residual(self.nModules) for _ in range(3)] if i == (self.depth - 1): res.append(self._make_residual(self.nModules)) # extra one for the middle hg.append(nn.ModuleList(res)) return nn.ModuleList(hg) def _hourglass_forward(self, depth_id, x): up_1 = self.hg[depth_id][0](x) low_1 = self.downsample(x) low_1 = self.hg[depth_id][1](low_1) if depth_id == (self.depth - 1): low_2 = self.hg[depth_id][3](low_1) else: low_2 = self._hourglass_forward(depth_id+1, low_1) low_3 = self.hg[depth_id][2](low_2) up_2 = self.upsample(low_3) return up_1 + up_2 def forward(self, x): return self._hourglass_forward(0, x) -
定义 Stacked Hourglass Network：
import torch.nn as nn from Model.HgResBlock import HgResBlock from Model.SingleHourglass import Hourglass class HourglassNet(nn.Module): def __init__(self, nStacks, nModules, nFeat, nClasses, resBlock=HgResBlock, inplanes=3): super(HourglassNet, self).__init__() self.nStacks = nStacks self.nModules = nModules self.nFeat = nFeat self.nClasses = nClasses self.resBlock = resBlock self.inplanes = inplanes hg, res, fc, score, fc_, score_ = [], [], [], [], [], [] for i in range(nStacks): hg.append(Hourglass(depth=4, nFeat=nFeat, nModules=nModules, resBlocks=resBlock)) res.append(self._make_residual(nModules)) fc.append(self._make_fc(nFeat, nFeat)) score.append(nn.Conv2d(nFeat, nClasses, kernel_size=1)) if i < (nStacks - 1): fc_.append(nn.Conv2d(nFeat, nFeat, kernel_size=1)) score_.append(nn.Conv2d(nClasses, nFeat, kernel_size=1)) self.hg = nn.ModuleList(hg) self.res = nn.ModuleList(res) self.fc = nn.ModuleList(fc) self.score = nn.ModuleList(score) self.fc_ = nn.ModuleList(fc_) self.score_ = nn.ModuleList(score_) def _make_head(self): self.conv_1 = nn.Conv2d(self.inplanes, 64, kernel_size=7, stride=2, padding=3) self.bn_1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.res_1 = self.resBlock(64, 128) self.pool = nn.MaxPool2d(2, 2) self.res_2 = self.resBlock(128, 128) self.res_3 = self.resBlock(128, self.nFeat) def _make_residual(self, n): return nn.Sequential(*[self.resBlock(self.nFeat, self.nFeat) for _ in range(n)]) def _make_fc(self, inplanes, outplanes): return nn.Sequential( nn.Conv2d(inplanes, outplanes, kernel_size=1), nn.BatchNorm2d(outplanes), nn.ReLU(True)) def forward(self, x): # head x = self.conv_1(x) x = self.bn_1(x) x = self.relu(x) x = self.res_1(x) x = self.pool(x) x = self.res_2(x) x = self.res_3(x) out = [] for i in range(self.nStacks): y = self.hg[i](x) y = self.res[i](y) y = self.fc[i](y) score = self.score[i](y) out.append(score) if i < (self.nStacks - 1): fc_ = self.fc_[i](y) score_ = self.score_[i](score) x = x + fc_ + score_ return out
References：
[1] Stacked Hourglass Networks for Human Pose Estimation
[2] [hourglass pytorch 实现]
(https://blog.csdn.net/github_36923418/article/details/81030883)
