Learning to downsample module. Args: in_channels (int): Number of input channels. dw_channels (tuple[int]): Number of output channels of the first and the second depthwise conv (dwconv) layers. out_channels (int): Number of output channels of the whole
| 11 | |
| 12 | |
| 13 | class LearningToDownsample(nn.Module): |
| 14 | """Learning to downsample module. |
| 15 | |
| 16 | Args: |
| 17 | in_channels (int): Number of input channels. |
| 18 | dw_channels (tuple[int]): Number of output channels of the first and |
| 19 | the second depthwise conv (dwconv) layers. |
| 20 | out_channels (int): Number of output channels of the whole |
| 21 | 'learning to downsample' module. |
| 22 | conv_cfg (dict | None): Config of conv layers. Default: None |
| 23 | norm_cfg (dict | None): Config of norm layers. Default: |
| 24 | dict(type='BN') |
| 25 | act_cfg (dict): Config of activation layers. Default: |
| 26 | dict(type='ReLU') |
| 27 | dw_act_cfg (dict): In DepthwiseSeparableConvModule, activation config |
| 28 | of depthwise ConvModule. If it is 'default', it will be the same |
| 29 | as `act_cfg`. Default: None. |
| 30 | """ |
| 31 | |
| 32 | def __init__(self, |
| 33 | in_channels, |
| 34 | dw_channels, |
| 35 | out_channels, |
| 36 | conv_cfg=None, |
| 37 | norm_cfg=dict(type='BN'), |
| 38 | act_cfg=dict(type='ReLU'), |
| 39 | dw_act_cfg=None): |
| 40 | super(LearningToDownsample, self).__init__() |
| 41 | self.conv_cfg = conv_cfg |
| 42 | self.norm_cfg = norm_cfg |
| 43 | self.act_cfg = act_cfg |
| 44 | self.dw_act_cfg = dw_act_cfg |
| 45 | dw_channels1 = dw_channels[0] |
| 46 | dw_channels2 = dw_channels[1] |
| 47 | |
| 48 | self.conv = ConvModule( |
| 49 | in_channels, |
| 50 | dw_channels1, |
| 51 | 3, |
| 52 | stride=2, |
| 53 | padding=1, |
| 54 | conv_cfg=self.conv_cfg, |
| 55 | norm_cfg=self.norm_cfg, |
| 56 | act_cfg=self.act_cfg) |
| 57 | |
| 58 | self.dsconv1 = DepthwiseSeparableConvModule( |
| 59 | dw_channels1, |
| 60 | dw_channels2, |
| 61 | kernel_size=3, |
| 62 | stride=2, |
| 63 | padding=1, |
| 64 | norm_cfg=self.norm_cfg, |
| 65 | dw_act_cfg=self.dw_act_cfg) |
| 66 | |
| 67 | self.dsconv2 = DepthwiseSeparableConvModule( |
| 68 | dw_channels2, |
| 69 | out_channels, |
| 70 | kernel_size=3, |