hrnet paddle实现，完成input、stage，未完成fuse

import paddle
import numpy as np
from paddle import nn
from paddle.nn import initializer
from paddle.nn import functional as F

BN_MOMENTUM = 0.2


class PlaceHolder(nn.Layer):

    def __init__(self):
        super(PlaceHolder, self).__init__()

    def forward(self, inputs):
        return inputs


# 通用3x3卷积块
class HRNetCov3x3(nn.Layer):
    def __init__(self, input_channels, out_channels, stride=1, padding=0):
        super(HRNetCov3x3, self).__init__()

        self.conv = nn.Conv2D(input_channels, out_channels, kernel_size=3,
                              stride=stride, padding=padding)
        self.bn = nn.BatchNorm2D(out_channels, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU()

    def forward(self, inputs):
        x = self.conv(inputs)
        x = self.bn(x)
        x = self.relu(x)

        return x


# stem
class HRNetStem(nn.Layer):
    def __init__(self, input_channels, out_channels, ):
        super(HRNetStem, self).__init__()
        self.conv1 = HRNetCov3x3(input_channels, out_channels, stride=2, padding=1)
        self.conv2 = HRNetCov3x3(out_channels, out_channels, stride=2, padding=1)

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.conv2(x)

        return x


class HRNetInput(nn.Layer):

    def __init__(self, input_channels, out_channels, stage1_inchannels):
        super(HRNetInput, self).__init__()
        self.stem = HRNetStem(input_channels, out_channels)
        self.in_change_conv = nn.Conv2D(out_channels, stage1_inchannels, kernel_size=1,
                                        stride=1, bias_attr=False)
        self.in_change_bn = nn.BatchNorm2D(stage1_inchannels, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU()

    def forward(self, inputs):
        x = self.stem(inputs)
        x = self.in_change_conv(x)
        x = self.in_change_bn(x)
        x = self.relu(x)

        return x


# 普通block
class NormalBlock(nn.Layer):
    def __init__(self, input_channels, out_channels):
        super(NormalBlock, self).__init__()
        self.conv1 = HRNetCov3x3(input_channels, out_channels,
                                 stride=1, padding=1)
        self.conv2 = HRNetCov3x3(input_channels, out_channels,
                                 stride=1, padding=1)

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.conv2(x)

        return x


# 残差block
class ResidualBlock(nn.Layer):
    def __init__(self, input_channels, out_channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = HRNetCov3x3(input_channels, out_channels,
                                 stride=1, padding=1)
        self.conv2 = nn.Conv2D(out_channels, out_channels, kernel_size=3,
                               stride=1, padding=1)
        self.bn2 = nn.BatchNorm2D(out_channels, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU()

    def forward(self, inputs):
        residual = inputs
        x = self.conv1(inputs)
        x = self.conv2(x)
        x = self.bn2(x)
        x += residual
        x = self.relu(x)

        return x


# Sequential
# LayerList
class HRNetStage(nn.Layer):
    def __init__(self, stage_channels, block):
        super(HRNetStage, self).__init__()
        self.stage_channels = stage_channels
        self.stage_branch_num = len(stage_channels)
        self.block = block
        self.block_num = 4
        self.stage_layers = self.create_stage_layers()

    def create_stage_layers(self):
        tostage_layers = []
        for i in range(self.stage_branch_num):
            branch_layer = []  # 串行
            for j in range(self.block_num):
                branch_layer.append(self.block(self.stage_channels[i],
                                               self.stage_channels[i]))
            branch_layer = nn.Sequential(*branch_layer)
            tostage_layers.append(branch_layer)

        return nn.LayerList(tostage_layers)

    def forward(self, inputs):
        outs = []
        for i in range(len(inputs)):
            x = inputs[i]
            out = self.stage_layers[i](x)
            outs.append(out)

        return outs


class HRNetTrans(nn.Layer):
    def __init__(self, old_branch_channels, new_branch_channels):
        super(HRNetTrans, self).__init__()
        self.old_branch_channels = old_branch_channels
        self.new_branch_channels = new_branch_channels
        self.old_branch_num = len(old_branch_channels)
        self.new_branch_num = len(new_branch_channels)
        self.trans_layers = self.create_new_branch_trans_layers()

    def create_new_branch_trans_layers(self):
        # LayerList
        totrns_layers = []
        for i in range(self.old_branch_num):
            brach_trans = []
            for j in range(self.new_branch_num):
                layer = []
                input_channels = self.old_branch_channels[i]
                if i == j:
                    layer.append(PlaceHolder())
                elif i < j:
                    for k in range(j, -1):
                        # 通道对齐
                        layer.append(nn.Conv2D(in_channels=input_channels,
                                               out_channels=self.new_branch_channels[j],
                                               kernel_size=1, bias_attr=False))
                        layer.append(nn.Conv2D(in_channels=input_channels,
                                               out_channels=self.new_branch_channels[j],
                                               kernel_size=1, stride=2, padding=1, bias_attr=False))
                        layer.append(nn.BatchNorm2D(self.new_branch_channels[j], momentum=BN_MOMENTUM))
                        layer.append(nn.ReLU())
                elif i > j:
                    for k in range(i-j):
                        # 通道对齐
                        layer.append(nn.Conv2D(in_channels=input_channels,
                                               out_channels=self.new_branch_channels[j],
                                               kernel_size=1, bias_attr=False))

                        layer.append(nn.BatchNorm2D(self.new_branch_channels[j], momentum=BN_MOMENTUM))
                        layer.append(nn.ReLU())
                        layer.append(nn.Upsample(scale_factor=2.))
                layer = nn.Sequential(*layer)
                brach_trans.append(layer)
            brach_trans = nn.LayerList(brach_trans)
            totrns_layers.append(brach_trans)

        return nn.LayerList(totrns_layers)

    def forward(self, inputs):
        outs = []
        for i in range(self.old_branch_num):
            x = inputs[i]
            out = []
            for j in range(self.new_branch_num):
                y = self.trans_layers[i][j]
                out.append(y)
            if len(outs) == 0:
                outs = out
            else:
                for i in range(self.new_branch_num):
                    outs[i] += out[i]
        return outs


class TestNet(nn.Layer):
    def __init__(self):
        super(TestNet, self).__init__()
        self.input = HRNetInput(3, out_channels=64, stage1_inchannels=32)
        self.stage1 = HRNetStage([32], NormalBlock)

        self.trans1 = HRNetTrans([32], [32, 64])

    def forward(self, inputs):
        x = self.input(inputs)
        x = [x]
        x = self.stage1(x)
        x = self.trans1(x)

        return x


if __name__ == "__main__":
    model = TestNet()
    data = np.random.randint(0, 256, (1, 3, 256, 256)).astype(np.float32)
    data = paddle.to_tensor(data)
    y = model(data)

    for i in range(len(y)):
       print(y[i])