| 289 | |
| 290 | # Spatial attention block (SAB) |
| 291 | class SAB(nn.Module): |
| 292 | def __init__(self, kernel_size=7): |
| 293 | super(SAB, self).__init__() |
| 294 | |
| 295 | assert kernel_size in (3, 7, 11), 'kernel must be 3 or 7 or 11' |
| 296 | padding = kernel_size//2 |
| 297 | |
| 298 | self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) |
| 299 | |
| 300 | self.sigmoid = nn.Sigmoid() |
| 301 | |
| 302 | self.init_weights('normal') |
| 303 | |
| 304 | def init_weights(self, scheme=''): |
| 305 | named_apply(partial(_init_weights, scheme=scheme), self) |
| 306 | |
| 307 | def forward(self, x): |
| 308 | avg_out = torch.mean(x, dim=1, keepdim=True) |
| 309 | max_out, _ = torch.max(x, dim=1, keepdim=True) |
| 310 | x = torch.cat([avg_out, max_out], dim=1) |
| 311 | x = self.conv(x) |
| 312 | return self.sigmoid(x) |
| 313 | |
| 314 | # Efficient multi-scale convolutional attention decoding (EMCAD) |
| 315 | class EMCAD(nn.Module): |