| import torch |
| import torch.nn as nn |
| from annotator.mmpkg.mmcv import is_tuple_of |
| from annotator.mmpkg.mmcv.cnn import ConvModule |
|
|
| from annotator.mmpkg.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) |
|
|
