Table of Contents
Shortcuts

# Source code for mmseg.models.decode_heads.nl_head

    import torch
    from mmcv.cnn import NonLocal2d

    from ..builder import HEADS
    from .fcn_head import FCNHead

    @HEADS.register_module()
    class NLHead(FCNHead):
        """Non-local Neural Networks.

        This head is the implementation of `NLNet
        <https://arxiv.org/abs/1711.07971>`_.

        Args:
            reduction (int): Reduction factor of projection transform. Default: 2.
            use_scale (bool): Whether to scale pairwise_weight by
                sqrt(1/inter_channels). Default: True.
            mode (str): The nonlocal mode. Options are 'embedded_gaussian',
                'dot_product'. Default: 'embedded_gaussian.'.
        """

        def __init__(self,
                     reduction=2,
                     use_scale=True,
                     mode='embedded_gaussian',
                     **kwargs):
            super(NLHead, self).__init__(num_convs=2, **kwargs)
            self.reduction = reduction
            self.use_scale = use_scale
            self.mode = mode
            self.nl_block = NonLocal2d(
                in_channels=self.channels,
                reduction=self.reduction,
                use_scale=self.use_scale,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                mode=self.mode)

        def forward(self, inputs):
            """Forward function."""
            x = self._transform_inputs(inputs)
            output = self.convs[0](x)
            output = self.nl_block(output)
            output = self.convs[1](output)
            if self.concat_input:
                output = self.conv_cat(torch.cat([x, output], dim=1))
            output = self.cls_seg(output)
            return output