OpenPose MobileNet V1 作为Backbone的模型结构（3）

OpenPose MobileNet V1 作为Backbone的模型结构（3）

Branch的合并

论文地址
论文中提到的源码地址
深度可分离卷积是如何减少计算量的
OpenPose MobileNet V1 作为Backbone的模型结构（1） MobileNet V1 的改造
OpenPose MobileNet V1 作为Backbone的模型结构（2） Initial Stage和Refinement Stage的优化
OpenPose MobileNet V1 作为Backbone的模型结构（3） Branch的合并
flyfish
论文原文
To produce new estimation of keypoint heatmaps and pafs the refinement stage takes features from backbone concatenated with previous estimation of keypoint heatmaps and pafs. Motivated by this fact we decided to share the most of computations between heatmaps and pafs and use single prediction branch in initial and refinement stage. We share all layers except the two last, which directly produce keypoint heatmaps and pafs.

1、上图左：VGG的OpenPose有两个prediction 分支，两个branch的结构是一样的，一个用来预测keypoint heatmap,一个用来预测PAFs。
2、上图右：MobileNet的OpenPose将两个branch合并成一个，只输出阶段使用1*1的卷积分出两个branch出来。一个用来预测keypoint heatmap,一个用来预测PAFs。

Initiate stage和Refinement stage这两个Stage各自都有合并
对于卷积通常封装成各种类型的卷积
例如
conv后面是否带BatchNorm2d，是否带ReLU
conv_dw深度可分离卷积后面是否BatchNorm2d，带ReLU还是带ELU

Initiate stage的branch的合并，下面的trunk是共享的

class InitialStage(nn.Module):
    def __init__(self, num_channels, num_heatmaps, num_pafs):
        super().__init__()
        self.trunk = nn.Sequential(
            conv(num_channels, num_channels, bn=False),
            conv(num_channels, num_channels, bn=False),
            conv(num_channels, num_channels, bn=False)
        )
        self.heatmaps = nn.Sequential(
            conv(num_channels, 512, kernel_size=1, padding=0, bn=False),
            conv(512, num_heatmaps, kernel_size=1, padding=0, bn=False, relu=False)
        )
        self.pafs = nn.Sequential(
            conv(num_channels, 512, kernel_size=1, padding=0, bn=False),
            conv(512, num_pafs, kernel_size=1, padding=0, bn=False, relu=False)
        )

    def forward(self, x):
        trunk_features = self.trunk(x)
        heatmaps = self.heatmaps(trunk_features)
        pafs = self.pafs(trunk_features)
        return [heatmaps, pafs]

Refinement stage的branch的合并，下面的trunk是共享的

对于Refinement stage先有Refinement stage block，多个这样的block组成Refinement stage

class RefinementStageBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.initial = conv(in_channels, out_channels, kernel_size=1, padding=0, bn=False)
        self.trunk = nn.Sequential(
            conv(out_channels, out_channels),
            conv(out_channels, out_channels, dilation=2, padding=2)
        )

    def forward(self, x):
        initial_features = self.initial(x)
        trunk_features = self.trunk(initial_features)
        return initial_features + trunk_features


class RefinementStage(nn.Module):
    def __init__(self, in_channels, out_channels, num_heatmaps, num_pafs):
        super().__init__()
        self.trunk = nn.Sequential(
            RefinementStageBlock(in_channels, out_channels),
            RefinementStageBlock(out_channels, out_channels),
            RefinementStageBlock(out_channels, out_channels),
            RefinementStageBlock(out_channels, out_channels),
            RefinementStageBlock(out_channels, out_channels)
        )
        self.heatmaps = nn.Sequential(
            conv(out_channels, out_channels, kernel_size=1, padding=0, bn=False),
            conv(out_channels, num_heatmaps, kernel_size=1, padding=0, bn=False, relu=False)
        )
        self.pafs = nn.Sequential(
            conv(out_channels, out_channels, kernel_size=1, padding=0, bn=False),
            conv(out_channels, num_pafs, kernel_size=1, padding=0, bn=False, relu=False)
        )

    def forward(self, x):
        trunk_features = self.trunk(x)
        heatmaps = self.heatmaps(trunk_features)
        pafs = self.pafs(trunk_features)
        return [heatmaps, pafs]