|
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667 |
- # Copyright (c) OpenMMLab. All rights reserved.
- import torch.nn.functional as F
- from mmcv.cnn import ConvModule
- from mmcv.runner 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
|
