| 28 | """ |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | in_channels, |
| 33 | out_channels, |
| 34 | stride, |
| 35 | norm_cfg=None, |
| 36 | act_cfg=None, |
| 37 | num_convs=4, |
| 38 | fusion_type="add", |
| 39 | init_cfg=None, |
| 40 | ): |
| 41 | super(STDCModule, self).__init__(init_cfg=init_cfg) |
| 42 | assert num_convs > 1 |
| 43 | assert fusion_type in ["add", "cat"] |
| 44 | self.stride = stride |
| 45 | self.with_downsample = True if self.stride == 2 else False |
| 46 | self.fusion_type = fusion_type |
| 47 | |
| 48 | self.layers = ModuleList() |
| 49 | conv_0 = ConvModule(in_channels, out_channels // 2, kernel_size=1, norm_cfg=norm_cfg) |
| 50 | |
| 51 | if self.with_downsample: |
| 52 | self.downsample = ConvModule( |
| 53 | out_channels // 2, |
| 54 | out_channels // 2, |
| 55 | kernel_size=3, |
| 56 | stride=2, |
| 57 | padding=1, |
| 58 | groups=out_channels // 2, |
| 59 | norm_cfg=norm_cfg, |
| 60 | act_cfg=None, |
| 61 | ) |
| 62 | |
| 63 | if self.fusion_type == "add": |
| 64 | self.layers.append(nn.Sequential(conv_0, self.downsample)) |
| 65 | self.skip = Sequential( |
| 66 | ConvModule( |
| 67 | in_channels, |
| 68 | in_channels, |
| 69 | kernel_size=3, |
| 70 | stride=2, |
| 71 | padding=1, |
| 72 | groups=in_channels, |
| 73 | norm_cfg=norm_cfg, |
| 74 | act_cfg=None, |
| 75 | ), |
| 76 | ConvModule(in_channels, out_channels, 1, norm_cfg=norm_cfg, act_cfg=None), |
| 77 | ) |
| 78 | else: |
| 79 | self.layers.append(conv_0) |
| 80 | self.skip = nn.AvgPool2d(kernel_size=3, stride=2, padding=1) |
| 81 | else: |
| 82 | self.layers.append(conv_0) |
| 83 | |
| 84 | for i in range(1, num_convs): |
| 85 | out_factor = 2 ** (i + 1) if i != num_convs - 1 else 2**i |
| 86 | self.layers.append( |
| 87 | ConvModule( |