Stacked Hourglass Networks for Human Pose Estimation 代码详解

论文整体简介

这是2016年中获得人体姿态识别得分最高的论文，在FLIC和MPII数据集上达到了state-of-the-art的效果。最近几年人体姿态估计（单人或多人）的研究几乎都是基于两种基本网络结构，一种是这篇论文的stacked hourglass模型，另一种是openPose论文的方法。

上图是论文提出的网络模型，叫做堆叠漏斗神经网络。这个网络架构形态就像它的名字的一样，是由一个个的漏斗状的神经网络级联起来，每一个漏斗神经网络就像编码器和解码器合成，负责提取特征和生成热图结果。整个网络使用了大量的卷积/反卷积层，池化/反池化层，ResNet以及全连接层。网络的输入是一张或者batchsize的标准大小尺寸图片（256×256），输出是该张图片缩小到一定尺寸的各个节点的热图（64×64）。

网络架构实现

作者最先是在torch上实现并开源了这篇文章的代码，地址：https://github.com/umich-vl/pose-hg-train ，这是训练部分的代码，demo部分的代码在作者另一个工程pose-hg-demo中。我们重点来看它的网络架构实现的部分，参数的设置和以及训练的数据的读取和预处理读者自己在代码中学习。
读者如果不愿意看枯燥的代码，我会一并贴出形象的结构图片，可参考图片理解。

创建网络

把打开工程文件夹，Stacked Hourglass的整体网络架构是在src/models/hg.lua文件中实现的：

function createModel()
    local inp = nn.Identity()()
    -- Initial processing of the image
    local cnv1_ = nnlib.SpatialConvolution(3,64,7,7,2,2,3,3)(inp)           -- 128
    local cnv1 = nnlib.ReLU(true)(nn.SpatialBatchNormalization(64)(cnv1_))
    local r1 = Residual(64,128)(cnv1)
    local pool = nnlib.SpatialMaxPooling(2,2,2,2)(r1)                       -- 64
    local r4 = Residual(128,128)(pool)
    local r5 = Residual(128,opt.nFeats)(r4)

    local out = {}
    local inter = r5

    for i = 1,opt.nStack do
        local hg = hourglass(4,opt.nFeats,inter)
        -- Residual layers at output resolution
        local ll = hg
        for j = 1,opt.nModules do ll = Residual(opt.nFeats,opt.nFeats)(ll) end
        -- Linear layer to produce first set of predictions
        ll = lin(opt.nFeats,opt.nFeats,ll)

        -- Predicted heatmaps
        local tmpOut = nnlib.SpatialConvolution(opt.nFeats,ref.nOutChannels,1,1,1,1,0,0)(ll)
        table.insert(out,tmpOut)

        -- Add predictions back
        if i < opt.nStack then
            local ll_ = nnlib.SpatialConvolution(opt.nFeats,opt.nFeats,1,1,1,1,0,0)(ll)
            local tmpOut_ = nnlib.SpatialConvolution(ref.nOutChannels,opt.nFeats,1,1,1,1,0,0)(tmpOut)
            inter = nn.CAddTable()({inter, ll_, tmpOut_})
        end
    end

    -- Final model
    local model = nn.gModule({inp}, out)

    return model

end

我们来一行行分析这个网络的实现。
inp是是表示输入的向量，在该工程中，它的大小为batchsize×3×256×256，训练的样本都是彩色图片，所以每一次训练的输入是batchsize大小的3通道的256×256的图片，如果图片本来不是256×256，要先截取人物部分的方框并缩放到256×256的大小。(以下输入输出大小均省去batchsize)
输入的图片先通过一个卷积层，卷积核大小为7*7，slide为2，pad为3，输出层数为64.经过这一层的卷积后，cnv1的大小为64*128*128，再把它batch normalization和用ReLU激活。
cnv1接下来通过一个Residual的网络，接下来的网络依次是一个池化层和两个Residual模块网络，到这边为止，得到将是opt.nFeats*64*64的特征，这里的opt.nFeats是预先设定的参数，论文和工程都是预设为256（后面opt开头的都是工程预设的数值）。

Residual模块

预处理和后面hourglass部分都用到了大量的Residual模块，其实现代码如下：

local conv = nnlib.SpatialConvolution
local batchnorm = nn.SpatialBatchNormalization
local relu = nnlib.ReLU

-- Main convolutional block
local function convBlock(numIn,numOut)
    return nn.Sequential()
        :add(batchnorm(numIn))
        :add(relu(true))
        :add(conv(numIn,numOut/2,1,1))
        :add(batchnorm(numOut/2))
        :add(relu(true))
        :add(conv(numOut/2,numOut/2,3,3,1,1,1,1))
        :add(batchnorm(numOut/2))
        :add(relu(true))
        :add(conv(numOut/2,numOut,1,1))
end

-- Skip layer
local function skipLayer(numIn,numOut)
    if numIn == numOut then
        return nn.Identity()
    else
        return nn.Sequential()
            :add(conv(numIn,numOut,1,1))
    end
end

-- Residual block
function Residual(numIn,numOut)
    return nn.Sequential()
        :add(nn.ConcatTable()
            :add(convBlock(numIn,numOut))
            :add(skipLayer(numIn,numOut)))
        :add(nn.CAddTable(true))
end

Residual模块结构如下图所示：

Residual模块是由两个子模块convBlock和skipLayer构成的，参数是输入层数和输出层数。
convBlock子模块相对复杂，由三组batchNormalization+ReLU+convolution串联构成，具体代码参考上面function convBlock部分，不再赘述。（图中的实线箭头移动部分，卷积核的大小如图中数字所示）skipLayer比较简单，如果输入层数等于输出层数，就直接输出，如果不等，就通过一个卷积层让输出层数变成设定值。（图中的虚线部分）最后两个子模块的输出合在一起，作为Residual模块的输出。

Houglass模块

从输入图片到得到opt.nFeats*64*64的特征，只是网络处理的开始，还没有真正进入漏斗状的网络中。接下来才真正要进入hourglass网络。
在for语句那一行的opt.nStack是预设的层叠个数，表示要叠用几个漏斗状网络。
每个层叠的漏斗网络包括基本的hourglass模块，opt.nModules个Residual模块，一个卷积层，batch normalization，ReLU以及再一个卷积层。到这里为止，网络得到就是nJoints*64*64大小的热图，这里的nJoints表示人体关键点的个数，也就是ref.nOutChannels的值。
基本的hourlass模块的函数如下：

local function hourglass(n, f, inp)
    -- Upper branch
    local up1 = inp
    for i = 1,opt.nModules do up1 = Residual(f,f)(up1) end

    -- Lower branch
    local low1 = nnlib.SpatialMaxPooling(2,2,2,2)(inp)
    for i = 1,opt.nModules do low1 = Residual(f,f)(low1) end
    local low2

    if n > 1 then low2 = hourglass(n-1,f,low1)
    else
        low2 = low1
        for i = 1,opt.nModules do low2 = Residual(f,f)(low2) end
    end

    local low3 = low2
    for i = 1,opt.nModules do low3 = Residual(f,f)(low3) end
    local up2 = nn.SpatialUpSamplingNearest(2)(low3)

    -- Bring two branches together
    return nn.CAddTable()({up1,up2})
end

hourglass模块的结构图如下：

从图中可以看出，这个hourglass模块具有对称的结构，从从中间依次扩展到两边，可以看做是小漏斗变成大漏斗，所以程序用递归的方式实现这个模块，图中的每个小块都可以看作是经过Residual模块和池化或反池化后的结果。而且前半部分的特征图会加到后半部分对称的位置。

Hourglass模块之间

hourlass模块的输出再经过两个1*1的卷积层后，结果作为当前stack的热图输出tempOut，两个卷积层之间的结果在createModel函数中记为ll。
在两个漏斗网络之间，还要对热图进行进一步的处理，分别是：1.对上一个漏斗网络中的ll通过一个卷积层；2.对上一个漏斗网络的热图tempOut通过一个卷积层。最后这两个结果和前一个漏斗的输入合并，作为下一个漏斗网络的输入。
漏斗间的处理结构图如下：

图中的蓝色部分表示当前stack的热图。

损失函数

最后得到的热图总共有opt.nStack组，每一个漏斗网络输出的结果都是一组热图（nJoints*64*64）。训练的时候，每一个漏斗网络的ground truth heatmaps都是一样的，根据关键点label生成一个二维的高斯图，损失函数是所有stack的输出热图和ground truth热图之差的L2范数（MSE）。

预测

预测时，输入任意一张人物图片（256*256),得到最后一个stack的输出热图（不是所有的）。算出每一个joint对应的热图中最大值元素所在的坐标，然后scale回去，作为最后预测到的关键点位置。

实战

作者在github上工程里，默认是设置8个stack，训练出来的模型在torch上达到205M，我后来对另一个tensorflow版本的代码进行修改，并按照与torch版本一样的方法进行训练，得到的模型达到301M，可以看出，该网络的参数是非常庞大的，肯定会有很多冗余的参数，如果想在移动端使用，需要进行进一步的修改和压缩。