PSPNet paddle实现

import numpy as np
import paddle.fluid as fluid
from paddle.fluid.dygraph import to_variable
from paddle.fluid.dygraph import Layer
from paddle.fluid.dygraph import Conv2D
from paddle.fluid.dygraph import BatchNorm
from paddle.fluid.dygraph import Dropout
from resnet_dilated import ResNet50

# pool with different bin_size
# interpolate back to input size
# concat
class PSPModule(Layer):
    def __init__(self, num_channels, box_list):
        super(PSPModule,self).__init__()
        self.num_channels = num_channels
        num_sub_channel = num_channels // len(box_list)
        self.box_list = box_list
        self.features = []
        for i in range(len(box_list)):
            self.features.append(
                fluid.dygraph.Sequential(
                    Conv2D(num_channels, num_sub_channel,1,padding=0,stride=1),
                    BatchNorm(num_sub_channel)
                )
            )



    def forward(self, inputs):
        out = [inputs]
        for i, out_size in enumerate(self.box_list):
            x = fluid.layers.adaptive_pool2d(inputs, out_size)
            x = self.features[i](x)
            x = fluid.layers.interpolate(x, inputs.shape[2::])
            out.append(x)
        out = fluid.layers.concat(out, axis=1)

        return out


class PSPNet(Layer):
    def __init__(self, num_classes=59, backbone='resnet50'):
        super(PSPNet, self).__init__()
        res = ResNet50()
        # stem: res.conv, res.pool2d_max
        self.conv = res.conv
        self.pool2d_max = res.pool2d_max
        self.layer1 = res.layer1
        self.layer2 = res.layer2
        self.layer3 = res.layer3
        self.layer4 = res.layer4

        box_list = [1,2,3,6]
        num_channels = 2048
        # psp: 2048 -> 2048*2
        self.pspmodule = PSPModule(num_channels,box_list)
        num_channels *= 2

        # cls: 2048*2 -> 512 -> num_classes
        self.cls = fluid.dygraph.Sequential(
            Conv2D(num_channels, 512,1,padding=0,stride=1),
            BatchNorm(512),
            Conv2D(512, num_classes,1,padding=0,stride=1),
            BatchNorm(num_classes)
        )
        # aux: 1024 -> 256 -> num_classes
        
    def forward(self, inputs):
        x = self.conv(inputs)
        x = self.pool2d_max(x)

        #print(x.shape)
        x = self.layer1(x)
        #print(x.shape)
        x = self.layer2(x)
        # print(x.shape)
        x = self.layer3(x)
        # print(x.shape)
        x = self.layer4(x)
        # print(x.shape)  [2, 2048, 60, 60]
        
        x = self.pspmodule(x)
        x = self.cls(x)
        x = fluid.layers.interpolate(x, inputs.shape[2::])
        # aux: tmp_x = layer3

        return x
            



def main():
    with fluid.dygraph.guard(fluid.CPUPlace()):
        x_data=np.random.rand(2,3, 473, 473).astype(np.float32)
        x = to_variable(x_data)
        model = PSPNet(num_classes=59)
        model.train()
        pred = model(x)
        print(pred.shape)

if __name__ =="__main__":
    main()