Table of Contents
Shortcuts

# Source code for mmseg.models.decode_heads.lraspp_head

    import torch
    import torch.nn as nn
    from mmcv import is_tuple_of
    from mmcv.cnn import ConvModule

    from mmseg.ops import resize
    from ..builder import HEADS
    from .decode_head import BaseDecodeHead

    @HEADS.register_module()
    class LRASPPHead(BaseDecodeHead):
        """Lite R-ASPP (LRASPP) head is proposed in Searching for MobileNetV3.

        This head is the improved implementation of `Searching for MobileNetV3
        <https://ieeexplore.ieee.org/document/9008835>`_.

        Args:
            branch_channels (tuple[int]): The number of output channels in every
                each branch. Default: (32, 64).
        """

        def __init__(self, branch_channels=(32, 64), **kwargs):
            super(LRASPPHead, self).__init__(**kwargs)
            if self.input_transform != 'multiple_select':
                raise ValueError('in Lite R-ASPP (LRASPP) head, input_transform '
                                 f'must be \'multiple_select\'. But received '
                                 f'\'{self.input_transform}\'')
            assert is_tuple_of(branch_channels, int)
            assert len(branch_channels) == len(self.in_channels) - 1
            self.branch_channels = branch_channels

            self.convs = nn.Sequential()
            self.conv_ups = nn.Sequential()
            for i in range(len(branch_channels)):
                self.convs.add_module(
                    f'conv{i}',
                    nn.Conv2d(
                        self.in_channels[i], branch_channels[i], 1, bias=False))
                self.conv_ups.add_module(
                    f'conv_up{i}',
                    ConvModule(
                        self.channels + branch_channels[i],
                        self.channels,
                        1,
                        norm_cfg=self.norm_cfg,
                        act_cfg=self.act_cfg,
                        bias=False))

            self.conv_up_input = nn.Conv2d(self.channels, self.channels, 1)

            self.aspp_conv = ConvModule(
                self.in_channels[-1],
                self.channels,
                1,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg,
                bias=False)
            self.image_pool = nn.Sequential(
                nn.AvgPool2d(kernel_size=49, stride=(16, 20)),
                ConvModule(
                    self.in_channels[2],
                    self.channels,
                    1,
                    act_cfg=dict(type='Sigmoid'),
                    bias=False))

        def forward(self, inputs):
            """Forward function."""
            inputs = self._transform_inputs(inputs)

            x = inputs[-1]

            x = self.aspp_conv(x) * resize(
                self.image_pool(x),
                size=x.size()[2:],
                mode='bilinear',
                align_corners=self.align_corners)
            x = self.conv_up_input(x)

            for i in range(len(self.branch_channels) - 1, -1, -1):
                x = resize(
                    x,
                    size=inputs[i].size()[2:],
                    mode='bilinear',
                    align_corners=self.align_corners)
                x = torch.cat([x, self.convs[i](inputs[i])], 1)
                x = self.conv_ups[i](x)

            return self.cls_seg(x)