| | |
| | import torch.nn.functional as F |
| | from mmcv.cnn import ConvModule |
| | from mmengine.model import BaseModule, ModuleList |
| |
|
| |
|
| | class ConvUpsample(BaseModule): |
| | """ConvUpsample performs 2x upsampling after Conv. |
| | |
| | There are several `ConvModule` layers. In the first few layers, upsampling |
| | will be applied after each layer of convolution. The number of upsampling |
| | must be no more than the number of ConvModule layers. |
| | |
| | Args: |
| | in_channels (int): Number of channels in the input feature map. |
| | inner_channels (int): Number of channels produced by the convolution. |
| | num_layers (int): Number of convolution layers. |
| | num_upsample (int | optional): Number of upsampling layer. Must be no |
| | more than num_layers. Upsampling will be applied after the first |
| | ``num_upsample`` layers of convolution. Default: ``num_layers``. |
| | conv_cfg (dict): Config dict for convolution layer. Default: None, |
| | which means using conv2d. |
| | norm_cfg (dict): Config dict for normalization layer. Default: None. |
| | init_cfg (dict): Config dict for initialization. Default: None. |
| | kwargs (key word augments): Other augments used in ConvModule. |
| | """ |
| |
|
| | def __init__(self, |
| | in_channels, |
| | inner_channels, |
| | num_layers=1, |
| | num_upsample=None, |
| | conv_cfg=None, |
| | norm_cfg=None, |
| | init_cfg=None, |
| | **kwargs): |
| | super(ConvUpsample, self).__init__(init_cfg) |
| | if num_upsample is None: |
| | num_upsample = num_layers |
| | assert num_upsample <= num_layers, \ |
| | f'num_upsample({num_upsample})must be no more than ' \ |
| | f'num_layers({num_layers})' |
| | self.num_layers = num_layers |
| | self.num_upsample = num_upsample |
| | self.conv = ModuleList() |
| | for i in range(num_layers): |
| | self.conv.append( |
| | ConvModule( |
| | in_channels, |
| | inner_channels, |
| | 3, |
| | padding=1, |
| | stride=1, |
| | conv_cfg=conv_cfg, |
| | norm_cfg=norm_cfg, |
| | **kwargs)) |
| | in_channels = inner_channels |
| |
|
| | def forward(self, x): |
| | num_upsample = self.num_upsample |
| | for i in range(self.num_layers): |
| | x = self.conv[i](x) |
| | if num_upsample > 0: |
| | num_upsample -= 1 |
| | x = F.interpolate( |
| | x, scale_factor=2, mode='bilinear', align_corners=False) |
| | return x |
| |
|
