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