pspnet.py

# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os

import paddle.nn.functional as F
from paddle import fluid
from paddle.fluid.dygraph import Conv2D

from dygraph.cvlibs import manager
from dygraph.models import model_utils
from dygraph.models.architectures import layer_utils
from dygraph.utils import utils


class PSPNet(fluid.dygraph.Layer):
    """
    The PSPNet implementation

    The orginal artile refers to 
        Zhao, Hengshuang, et al. "Pyramid scene parsing network." 
        Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.
        (https://openaccess.thecvf.com/content_cvpr_2017/papers/Zhao_Pyramid_Scene_Parsing_CVPR_2017_paper.pdf)

    Args:
        backbone (str): backbone name, currently support Resnet50/101.

        num_classes (int): the unique number of target classes. Default 2.

        output_stride (int): the ratio of input size and final feature size. Default 16.

        backbone_indices (tuple): two values in the tuple indicte the indices of output of backbone.
                        the first index will be taken as a deep-supervision feature in auxiliary layer;
                        the second one will be taken as input of Pyramid Pooling Module (PPModule).
                        Usually backbone consists of four downsampling stage, and return an output of
                        each stage, so we set default (2, 3), which means taking feature map of the third
                        stage (res4b22) in backbone, and feature map of the fourth stage (res5c) as input of PPModule.

        backbone_channels (tuple): the same length with "backbone_indices". It indicates the channels of corresponding index.

        pp_out_channels (int): output channels after Pyramid Pooling Module. Default to 1024.

        bin_sizes (tuple): the out size of pooled feature maps. Default to (1,2,3,6).

        enable_auxiliary_loss (bool): a bool values indictes whether adding auxiliary loss. Default to True.

        ignore_index (int): the value of ground-truth mask would be ignored while doing evaluation. Default to 255.

        pretrained_model (str): the pretrained_model path of backbone.
    """

    def __init__(self,
                 backbone,
                 num_classes=2,
                 output_stride=16,
                 backbone_indices=(2, 3),
                 backbone_channels=(1024, 2048),
                 pp_out_channels=1024,
                 bin_sizes=(1, 2, 3, 6),
                 enable_auxiliary_loss=True,
                 ignore_index=255,
                 pretrained_model=None):

        super(PSPNet, self).__init__()
        self.backbone = manager.BACKBONES[backbone](output_stride=output_stride,
                                                    multi_grid=(1, 1, 1))
        self.backbone_indices = backbone_indices

        self.psp_module = PPModule(in_channels=backbone_channels[1],
                                   out_channels=pp_out_channels,
                                   bin_sizes=bin_sizes)

        self.conv = Conv2D(num_channels=pp_out_channels,
                           num_filters=num_classes,
                           filter_size=1)

        if enable_auxiliary_loss:
            self.fcn_head = model_utils.FCNHead(in_channels=backbone_channels[0], out_channels=num_classes)

        self.enable_auxiliary_loss = enable_auxiliary_loss
        self.ignore_index = ignore_index

        self.init_weight(pretrained_model)

    def forward(self, input, label=None):

        _, feat_list = self.backbone(input)

        x = feat_list[self.backbone_indices[1]]
        x = self.psp_module(x)
        x = F.dropout(x, dropout_prob=0.1)
        logit = self.conv(x)
        logit = fluid.layers.resize_bilinear(logit, input.shape[2:])

        if self.enable_auxiliary_loss:
            auxiliary_feat = feat_list[self.backbone_indices[0]]
            auxiliary_logit = self.fcn_head(auxiliary_feat)
            auxiliary_logit = fluid.layers.resize_bilinear(auxiliary_logit, input.shape[2:])

        if self.training:
            loss = model_utils.get_loss(logit, label)
            if self.enable_auxiliary_loss:
                auxiliary_loss = model_utils.get_loss(auxiliary_logit, label)
                loss += (0.4 * auxiliary_loss)
            return loss


        else:
            pred, score_map = model_utils.get_pred_score_map(logit)
            return pred, score_map

    def init_weight(self, pretrained_model=None):
        """
        Initialize the parameters of model parts.

        Args:
            pretrained_model ([str], optional): the pretrained_model path of backbone. Defaults to None.
        """

        if pretrained_model is not None:
            if os.path.exists(pretrained_model):
                utils.load_pretrained_model(self.backbone, pretrained_model)


class PPModule(fluid.dygraph.Layer):
    """
    Pyramid pooling module

    Args:
        in_channels (int): the number of intput channels to pyramid pooling module.

        out_channels (int): the number of output channels after pyramid pooling module.

        bin_sizes (tuple): the out size of pooled feature maps. Default to (1,2,3,6).
    """

    def __init__(self, in_channels, out_channels, bin_sizes=(1, 2, 3, 6)):
        super(PPModule, self).__init__()
        self.bin_sizes = bin_sizes

        # we use dimension reduction after pooling mentioned in original implementation.
        self.stages = fluid.dygraph.LayerList([self._make_stage(in_channels, size) for size in bin_sizes])

        self.conv_bn_relu2 = layer_utils.ConvBnRelu(num_channels=in_channels * 2,
                                                    num_filters=out_channels,
                                                    filter_size=3,
                                                    padding=1)

    def _make_stage(self, in_channels, size):
        """
        Create one pooling layer.

        In our implementation, we adopt the same dimention reduction as the original paper that might be
        slightly different with other implementations. 

        After pooling, the channels are reduced to 1/len(bin_sizes) immediately, while some other implementations
        keep the channels to be same.


        Args:
            in_channels (int): the number of intput channels to pyramid pooling module.

            size (int): the out size of the pooled layer.

        Returns:
            conv (tensor): a tensor after Pyramid Pooling Module
        """

        # this paddle version does not support AdaptiveAvgPool2d, so skip it here.
        # prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = layer_utils.ConvBnRelu(num_channels=in_channels,
                                      num_filters=in_channels // len(self.bin_sizes),
                                      filter_size=1)

        return conv

    def forward(self, input):
        cat_layers = []
        for i, stage in enumerate(self.stages):
            size = self.bin_sizes[i]
            x = fluid.layers.adaptive_pool2d(input, pool_size=(size, size), pool_type="max")
            x = stage(x)
            x = fluid.layers.resize_bilinear(x, out_shape=input.shape[2:])
            cat_layers.append(x)
        cat_layers = [input] + cat_layers[::-1]
        cat = fluid.layers.concat(cat_layers, axis=1)
        out = self.conv_bn_relu2(cat)

        return out


@manager.MODELS.add_component
def pspnet_resnet101_vd(*args, **kwargs):
    pretrained_model = None
    return PSPNet(backbone='ResNet101_vd', pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet101_vd_os8(*args, **kwargs):
    pretrained_model = None
    return PSPNet(backbone='ResNet101_vd', output_stride=8, pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet50_vd(*args, **kwargs):
    pretrained_model = None
    return PSPNet(backbone='ResNet50_vd', pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet50_vd_os8(*args, **kwargs):
    pretrained_model = None
    return PSPNet(backbone='ResNet50_vd', output_stride=8, pretrained_model=pretrained_model, **kwargs)