HRNet网络结构及TensorFlow实现

HRNet网络结构

HRNet是王井东老师团队提出的网络框架，据王老师介绍，HRNet将是继AlexNet、GoogleNet、VGGNet、ResNet、DenseNet之后，用于多种视觉任务的新颖的网络框架。目前已经在人体姿态估计、关键点检测、图像分割、物体识别、物体分类等任务上取得了出色的结果。
论文地址（HRNet V1）：Deep High-Resolution Representation Learning for Human Pose Estimation
论文地址（HRNet V2）：Deep High-Resolution Representation Learning for Visual Recognition
开源地址（Pytorch）：https://github.com/HRNet
由于王老师团队使用Pytorch实现的，对于熟悉TensorFlow的同学了解起来比较吃力（虽然TensorFlow和Pytorch都应该掌握）。所以博主在查阅资料后收集了HRNet得到TensorFlow实现。
开源地址（TensorFlow）：https://github.com/AI-Chen/HRNet-V2
由于TensorFlow是有其他人复现的，所以与官方版本存在一定的出入，因此建议结合论文阅读代码，部分可能存在问题的程序建议结合Pytorch确认。博主发现了TensorFlow版本存在的一个问题，将在后面提出。
在这篇博客中主要想介绍HRNet的网络结构。由于HRNet网络结构比较庞大，所以得静下心来认真阅读源代码才能真正的了解HRNet的网络结构。首先是HRNet论文中的网络结构

以下是博主通过自己的理解绘制的网络结构，不适之处希望指出。

在网络结构中每一个conv之后都接了BN，而且在大部分conv之后都会使用Relu作为激活函数，由于篇幅原因，并没有在网络结构中体现。只有部分block（Basic Block, Bottleneck Block）是在求和之后接的Relu。

在绘制的网络结构中Fusion Layer可能比较杂乱无章，因此给出论文中的Fusion 模块的示意图。

为了适用于不同的任务，网络针对于主干网络的输出提出了不同的处理方式，即不同的head。

（a）HRNet V1(用于姿态估计) （b）HRNet V2(用于语义分割) （c）HRNet V2p （用于目标检测）另外还有一种用于分类任务的处理方式

网络结构结合论文以及程序一起看能够起到事半功倍的效果。
虽然在在论文中王老师一直强调HRNet能够维持高分辨率表示，而不是像之前的网络框架那样从低分辨率中恢复高分辨。但是仔细阅读程序实现会发现，为了减少GPU内存的使用以及降低计算量，在网络之初，已经进行了两次下采样，而最后只是粗略的使用一次resize来恢复原始的分辨率。
另外在TensorFlow实现的程序中，个人感觉 utils\HighResolutionModule.py 中的_make_fuse_layers是有问题。

def _make_fuse_layers(self,x):
        if self.num_branches == 1:
            return x
        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        fuse_layers = []
        for i in range(num_branches if self.multi_scale_output else 1):
            temp_layer = []
            for j in range(num_branches):
                if j > i:
                    with tf.variable_scope(self.stage + '_fuse_layers' + str(i)+str(j)):
                        temp = block_Model.conv(x[j], 1, num_inchannels[i], 'conv_', stride=1)
                        temp = block_Model.batch_norm(temp, self._is_training)
                        temp = tf.image.resize_bilinear(temp, x[i].get_shape().as_list()[1:3])
                        temp_layer.append(temp)
                elif j == i:
                    temp_layer.append(x[j])
                else:
                    with tf.variable_scope(self.stage + '_fuse_layers' + str(i) + str(j)):
                        temp = block_Model.conv(x[j], 3, num_inchannels[i], 'conv_', stride=2)
                        temp = block_Model.batch_norm(temp, self._is_training)
                        temp_layer.append(temp)
            for k in range(len(temp_layer)):
                fuse = 0
                fuse = fuse + temp_layer[k]
            fuse_layers.append(fuse)
        return fuse_layers

当j<i时，只进行了一次下采样，事实上应该根据i-j来确定下采样的次数。根据论文及Pytorch的实现重写了该方法。

    def _make_fuse_layers(self,x):
        if self.num_branches == 1:
            return x
        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        fuse_layers = []
        for i in range(num_branches if self.multi_scale_output else 1):
            temp_layer = []
            for j in range(num_branches):
                if j > i:
                    with tf.variable_scope(self.stage + '_fuse_layers' +  '_' + str(j)+ '_' + str(i)):
                        temp = block_Model.conv(x[j], 1, num_inchannels[i], 'conv_', stride=1)
                        temp = block_Model.batch_norm(temp, self._is_training)
                        temp = tf.image.resize_bilinear(temp, x[i].get_shape().as_list()[1:3])
                        temp_layer.append(temp)
                elif j == i:
                    temp_layer.append(x[j])
                else:
                    temp = x[j]
                    for k in range(i - j):
                        if k == i - j - 1:
                            with tf.variable_scope(self.stage + '_fuse_layers' + '_' + str(j) +  '_' + str(i) + '_' + str(k)):
                                temp = block_Model.conv(temp, 3, num_inchannels[i], 'conv_', stride=2)
                                temp = block_Model.batch_norm(temp, self._is_training)
                        else:
                            with tf.variable_scope(self.stage + '_fuse_layers' + '_' + str(j) +  '_' + str(i) + '_' + str(k)):
                                temp = block_Model.conv(temp, 3, num_inchannels[i], 'conv_', stride=2)
                                temp = block_Model.batch_norm(temp, self._is_training)
                                temp = tf.nn.relu(temp)
                    temp_layer.append(temp)
            for k in range(len(temp_layer)):
                print("stage{}, branche{}, temp{} shape:{}".format(num_branches, i, k, temp_layer[k].get_shape().as_list()))
                fuse = 0
                fuse = fuse + temp_layer[k]
            print("stage{}, branche{} fused shape: {}".format(num_branches, i, fuse.get_shape().as_list()))
            fuse_layers.append(fuse)
        return fuse_layers

个人感觉HRNet中通过多分辨率信息交互，能够离开一些高级语义特征增强高分辨率特征的语义表示，而且它在空间上也会更加精确。最终促进许多对位置比较敏感的任务。
以上只是博主个人的一些理解，才疏学浅，不适之处，希望大家指正。