| 178 | |
| 179 | |
| 180 | class _OSA_module(nn.Module): |
| 181 | def __init__( |
| 182 | self, in_ch, stage_ch, concat_ch, layer_per_block, module_name, SE=False, identity=False, depthwise=False |
| 183 | ): |
| 184 | |
| 185 | super(_OSA_module, self).__init__() |
| 186 | |
| 187 | self.identity = identity |
| 188 | self.depthwise = depthwise |
| 189 | self.isReduced = False |
| 190 | self.layers = nn.ModuleList() |
| 191 | in_channel = in_ch |
| 192 | if self.depthwise and in_channel != stage_ch: |
| 193 | self.isReduced = True |
| 194 | self.conv_reduction = nn.Sequential( |
| 195 | OrderedDict(conv1x1(in_channel, stage_ch, "{}_reduction".format(module_name), "0")) |
| 196 | ) |
| 197 | for i in range(layer_per_block): |
| 198 | if self.depthwise: |
| 199 | self.layers.append(nn.Sequential(OrderedDict(dw_conv3x3(stage_ch, stage_ch, module_name, i)))) |
| 200 | else: |
| 201 | self.layers.append(nn.Sequential(OrderedDict(conv3x3(in_channel, stage_ch, module_name, i)))) |
| 202 | in_channel = stage_ch |
| 203 | |
| 204 | # feature aggregation |
| 205 | in_channel = in_ch + layer_per_block * stage_ch |
| 206 | self.concat = nn.Sequential(OrderedDict(conv1x1(in_channel, concat_ch, module_name, "concat"))) |
| 207 | |
| 208 | self.ese = eSEModule(concat_ch) |
| 209 | |
| 210 | def forward(self, x): |
| 211 | |
| 212 | identity_feat = x |
| 213 | |
| 214 | output = [] |
| 215 | output.append(x) |
| 216 | if self.depthwise and self.isReduced: |
| 217 | x = self.conv_reduction(x) |
| 218 | for layer in self.layers: |
| 219 | x = layer(x) |
| 220 | output.append(x) |
| 221 | |
| 222 | x = torch.cat(output, dim=1) |
| 223 | xt = self.concat(x) |
| 224 | |
| 225 | xt = self.ese(xt) |
| 226 | |
| 227 | if self.identity: |
| 228 | xt = xt + identity_feat |
| 229 | |
| 230 | return xt |
| 231 | |
| 232 | |
| 233 | class _OSA_stage(nn.Sequential): |