| import torch |
| import torch.nn as nn |
| from mmcv.cnn import ConvModule |
|
|
| from ..builder import HEADS |
| from ..utils import SelfAttentionBlock as _SelfAttentionBlock |
| from .decode_head import BaseDecodeHead |
|
|
|
|
| class PPMConcat(nn.ModuleList): |
| """Pyramid Pooling Module that only concat the features of each layer. |
| |
| Args: |
| pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid |
| Module. |
| """ |
|
|
| def __init__(self, pool_scales=(1, 3, 6, 8)): |
| super(PPMConcat, self).__init__( |
| [nn.AdaptiveAvgPool2d(pool_scale) for pool_scale in pool_scales]) |
|
|
| def forward(self, feats): |
| """Forward function.""" |
| ppm_outs = [] |
| for ppm in self: |
| ppm_out = ppm(feats) |
| ppm_outs.append(ppm_out.view(*feats.shape[:2], -1)) |
| concat_outs = torch.cat(ppm_outs, dim=2) |
| return concat_outs |
|
|
|
|
| class SelfAttentionBlock(_SelfAttentionBlock): |
| """Make a ANN used SelfAttentionBlock. |
| |
| Args: |
| low_in_channels (int): Input channels of lower level feature, |
| which is the key feature for self-attention. |
| high_in_channels (int): Input channels of higher level feature, |
| which is the query feature for self-attention. |
| channels (int): Output channels of key/query transform. |
| out_channels (int): Output channels. |
| share_key_query (bool): Whether share projection weight between key |
| and query projection. |
| query_scale (int): The scale of query feature map. |
| key_pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid |
| Module of key feature. |
| conv_cfg (dict|None): Config of conv layers. |
| norm_cfg (dict|None): Config of norm layers. |
| act_cfg (dict|None): Config of activation layers. |
| """ |
|
|
| def __init__(self, low_in_channels, high_in_channels, channels, |
| out_channels, share_key_query, query_scale, key_pool_scales, |
| conv_cfg, norm_cfg, act_cfg): |
| key_psp = PPMConcat(key_pool_scales) |
| if query_scale > 1: |
| query_downsample = nn.MaxPool2d(kernel_size=query_scale) |
| else: |
| query_downsample = None |
| super(SelfAttentionBlock, self).__init__( |
| key_in_channels=low_in_channels, |
| query_in_channels=high_in_channels, |
| channels=channels, |
| out_channels=out_channels, |
| share_key_query=share_key_query, |
| query_downsample=query_downsample, |
| key_downsample=key_psp, |
| key_query_num_convs=1, |
| key_query_norm=True, |
| value_out_num_convs=1, |
| value_out_norm=False, |
| matmul_norm=True, |
| with_out=True, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
|
|
|
|
| class AFNB(nn.Module): |
| """Asymmetric Fusion Non-local Block(AFNB) |
| |
| Args: |
| low_in_channels (int): Input channels of lower level feature, |
| which is the key feature for self-attention. |
| high_in_channels (int): Input channels of higher level feature, |
| which is the query feature for self-attention. |
| channels (int): Output channels of key/query transform. |
| out_channels (int): Output channels. |
| and query projection. |
| query_scales (tuple[int]): The scales of query feature map. |
| Default: (1,) |
| key_pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid |
| Module of key feature. |
| conv_cfg (dict|None): Config of conv layers. |
| norm_cfg (dict|None): Config of norm layers. |
| act_cfg (dict|None): Config of activation layers. |
| """ |
|
|
| def __init__(self, low_in_channels, high_in_channels, channels, |
| out_channels, query_scales, key_pool_scales, conv_cfg, |
| norm_cfg, act_cfg): |
| super(AFNB, self).__init__() |
| self.stages = nn.ModuleList() |
| for query_scale in query_scales: |
| self.stages.append( |
| SelfAttentionBlock( |
| low_in_channels=low_in_channels, |
| high_in_channels=high_in_channels, |
| channels=channels, |
| out_channels=out_channels, |
| share_key_query=False, |
| query_scale=query_scale, |
| key_pool_scales=key_pool_scales, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg)) |
| self.bottleneck = ConvModule( |
| out_channels + high_in_channels, |
| out_channels, |
| 1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=None) |
|
|
| def forward(self, low_feats, high_feats): |
| """Forward function.""" |
| priors = [stage(high_feats, low_feats) for stage in self.stages] |
| context = torch.stack(priors, dim=0).sum(dim=0) |
| output = self.bottleneck(torch.cat([context, high_feats], 1)) |
| return output |
|
|
|
|
| class APNB(nn.Module): |
| """Asymmetric Pyramid Non-local Block (APNB) |
| |
| Args: |
| in_channels (int): Input channels of key/query feature, |
| which is the key feature for self-attention. |
| channels (int): Output channels of key/query transform. |
| out_channels (int): Output channels. |
| query_scales (tuple[int]): The scales of query feature map. |
| Default: (1,) |
| key_pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid |
| Module of key feature. |
| conv_cfg (dict|None): Config of conv layers. |
| norm_cfg (dict|None): Config of norm layers. |
| act_cfg (dict|None): Config of activation layers. |
| """ |
|
|
| def __init__(self, in_channels, channels, out_channels, query_scales, |
| key_pool_scales, conv_cfg, norm_cfg, act_cfg): |
| super(APNB, self).__init__() |
| self.stages = nn.ModuleList() |
| for query_scale in query_scales: |
| self.stages.append( |
| SelfAttentionBlock( |
| low_in_channels=in_channels, |
| high_in_channels=in_channels, |
| channels=channels, |
| out_channels=out_channels, |
| share_key_query=True, |
| query_scale=query_scale, |
| key_pool_scales=key_pool_scales, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg)) |
| self.bottleneck = ConvModule( |
| 2 * in_channels, |
| out_channels, |
| 1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
|
|
| def forward(self, feats): |
| """Forward function.""" |
| priors = [stage(feats, feats) for stage in self.stages] |
| context = torch.stack(priors, dim=0).sum(dim=0) |
| output = self.bottleneck(torch.cat([context, feats], 1)) |
| return output |
|
|
|
|
| @HEADS.register_module() |
| class ANNHead(BaseDecodeHead): |
| """Asymmetric Non-local Neural Networks for Semantic Segmentation. |
| |
| This head is the implementation of `ANNNet |
| <https://arxiv.org/abs/1908.07678>`_. |
| |
| Args: |
| project_channels (int): Projection channels for Nonlocal. |
| query_scales (tuple[int]): The scales of query feature map. |
| Default: (1,) |
| key_pool_scales (tuple[int]): The pooling scales of key feature map. |
| Default: (1, 3, 6, 8). |
| """ |
|
|
| def __init__(self, |
| project_channels, |
| query_scales=(1, ), |
| key_pool_scales=(1, 3, 6, 8), |
| **kwargs): |
| super(ANNHead, self).__init__( |
| input_transform='multiple_select', **kwargs) |
| assert len(self.in_channels) == 2 |
| low_in_channels, high_in_channels = self.in_channels |
| self.project_channels = project_channels |
| self.fusion = AFNB( |
| low_in_channels=low_in_channels, |
| high_in_channels=high_in_channels, |
| out_channels=high_in_channels, |
| channels=project_channels, |
| query_scales=query_scales, |
| key_pool_scales=key_pool_scales, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
| self.bottleneck = ConvModule( |
| high_in_channels, |
| self.channels, |
| 3, |
| padding=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
| self.context = APNB( |
| in_channels=self.channels, |
| out_channels=self.channels, |
| channels=project_channels, |
| query_scales=query_scales, |
| key_pool_scales=key_pool_scales, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
|
|
| def forward(self, inputs): |
| """Forward function.""" |
| low_feats, high_feats = self._transform_inputs(inputs) |
| output = self.fusion(low_feats, high_feats) |
| output = self.dropout(output) |
| output = self.bottleneck(output) |
| output = self.context(output) |
| output = self.cls_seg(output) |
|
|
| return output |
|
|
