| 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 |
|
|
