DeepLabV3Plus-based mscale segmentation model
| 230 | |
| 231 | |
| 232 | class MscaleV3Plus(MscaleBase): |
| 233 | """ |
| 234 | DeepLabV3Plus-based mscale segmentation model |
| 235 | """ |
| 236 | def __init__(self, num_classes, trunk='wrn38', criterion=None, |
| 237 | use_dpc=False, fuse_aspp=False, attn_2b=False): |
| 238 | super(MscaleV3Plus, self).__init__() |
| 239 | self.criterion = criterion |
| 240 | self.fuse_aspp = fuse_aspp |
| 241 | self.attn_2b = attn_2b |
| 242 | self.backbone, s2_ch, _s4_ch, high_level_ch = get_trunk(trunk) |
| 243 | self.aspp, aspp_out_ch = get_aspp(high_level_ch, |
| 244 | bottleneck_ch=256, |
| 245 | output_stride=8, |
| 246 | dpc=use_dpc) |
| 247 | self.bot_fine = nn.Conv2d(s2_ch, 48, kernel_size=1, bias=False) |
| 248 | self.bot_aspp = nn.Conv2d(aspp_out_ch, 256, kernel_size=1, bias=False) |
| 249 | |
| 250 | # Semantic segmentation prediction head |
| 251 | bot_ch = cfg.MODEL.SEGATTN_BOT_CH |
| 252 | self.final = nn.Sequential( |
| 253 | nn.Conv2d(256 + 48, bot_ch, kernel_size=3, padding=1, bias=False), |
| 254 | Norm2d(bot_ch), |
| 255 | nn.ReLU(inplace=True), |
| 256 | nn.Conv2d(bot_ch, bot_ch, kernel_size=3, padding=1, bias=False), |
| 257 | Norm2d(bot_ch), |
| 258 | nn.ReLU(inplace=True), |
| 259 | nn.Conv2d(bot_ch, num_classes, kernel_size=1, bias=False)) |
| 260 | |
| 261 | # Scale-attention prediction head |
| 262 | if self.attn_2b: |
| 263 | attn_ch = 2 |
| 264 | else: |
| 265 | attn_ch = 1 |
| 266 | |
| 267 | scale_in_ch = 256 + 48 |
| 268 | |
| 269 | self.scale_attn = make_attn_head(in_ch=scale_in_ch, |
| 270 | out_ch=attn_ch) |
| 271 | |
| 272 | if cfg.OPTIONS.INIT_DECODER: |
| 273 | initialize_weights(self.bot_fine) |
| 274 | initialize_weights(self.bot_aspp) |
| 275 | initialize_weights(self.scale_attn) |
| 276 | initialize_weights(self.final) |
| 277 | else: |
| 278 | initialize_weights(self.final) |
| 279 | |
| 280 | def _build_scale_tensor(self, scale_float, shape): |
| 281 | """ |
| 282 | Fill a 2D tensor with a constant scale value |
| 283 | """ |
| 284 | bs = scale_float.shape[0] |
| 285 | scale_tensor = None |
| 286 | for b in range(bs): |
| 287 | a_tensor = torch.Tensor(1, 1, *shape) |
| 288 | a_tensor.fill_(scale_float[b]) |
| 289 | if scale_tensor is None: |
no outgoing calls
no test coverage detected