pspnet.py 8.9 KB
Newer Older
M
michaelowenliu 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146
# 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).
M
michaelowenliu 已提交
147 148

        dim_reduction (bool): a bool value represent if reduing dimention after pooling. Default to True.
M
michaelowenliu 已提交
149 150
    """

M
michaelowenliu 已提交
151
    def __init__(self, in_channels, out_channels, bin_sizes=(1, 2, 3, 6), dim_reduction=True):
M
michaelowenliu 已提交
152 153 154
        super(PPModule, self).__init__()
        self.bin_sizes = bin_sizes

M
michaelowenliu 已提交
155 156 157 158
        inter_channels = in_channels
        if dim_reduction:
            inter_channels = in_channels // len(bin_sizes)
        
M
michaelowenliu 已提交
159
        # we use dimension reduction after pooling mentioned in original implementation.
M
michaelowenliu 已提交
160
        self.stages = fluid.dygraph.LayerList([self._make_stage(in_channels, inter_channels, size) for size in bin_sizes])
M
michaelowenliu 已提交
161

M
michaelowenliu 已提交
162
        self.conv_bn_relu2 = layer_utils.ConvBnRelu(num_channels=in_channels + inter_channels * len(bin_sizes),
M
michaelowenliu 已提交
163 164 165 166
                                                    num_filters=out_channels,
                                                    filter_size=3,
                                                    padding=1)

M
michaelowenliu 已提交
167
    def _make_stage(self, in_channels, out_channels, size):
M
michaelowenliu 已提交
168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
        """
        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,
M
michaelowenliu 已提交
190
                                      num_filters=out_channels,
M
michaelowenliu 已提交
191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
                                      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):
M
michaelowenliu 已提交
212
    pretrained_model = None
M
michaelowenliu 已提交
213 214 215 216 217
    return PSPNet(backbone='ResNet101_vd', pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet101_vd_os8(*args, **kwargs):
M
michaelowenliu 已提交
218
    pretrained_model = None
M
michaelowenliu 已提交
219 220 221 222 223
    return PSPNet(backbone='ResNet101_vd', output_stride=8, pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet50_vd(*args, **kwargs):
M
michaelowenliu 已提交
224
    pretrained_model = None
M
michaelowenliu 已提交
225 226 227 228 229
    return PSPNet(backbone='ResNet50_vd', pretrained_model=pretrained_model, **kwargs)


@manager.MODELS.add_component
def pspnet_resnet50_vd_os8(*args, **kwargs):
M
michaelowenliu 已提交
230
    pretrained_model = None
M
michaelowenliu 已提交
231
    return PSPNet(backbone='ResNet50_vd', output_stride=8, pretrained_model=pretrained_model, **kwargs)