| 19 | |
| 20 | |
| 21 | class SelectiveKernelAttn(nn.Module): |
| 22 | def __init__(self, channels, num_paths=2, attn_channels=32, |
| 23 | act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d): |
| 24 | """ Selective Kernel Attention Module |
| 25 | |
| 26 | Selective Kernel attention mechanism factored out into its own module. |
| 27 | |
| 28 | """ |
| 29 | super(SelectiveKernelAttn, self).__init__() |
| 30 | self.num_paths = num_paths |
| 31 | self.fc_reduce = nn.Conv2d(channels, attn_channels, kernel_size=1, bias=False) |
| 32 | self.bn = norm_layer(attn_channels) |
| 33 | self.act = act_layer(inplace=True) |
| 34 | self.fc_select = nn.Conv2d(attn_channels, channels * num_paths, kernel_size=1, bias=False) |
| 35 | |
| 36 | def forward(self, x): |
| 37 | assert x.shape[1] == self.num_paths |
| 38 | x = x.sum(1).mean((2, 3), keepdim=True) |
| 39 | x = self.fc_reduce(x) |
| 40 | x = self.bn(x) |
| 41 | x = self.act(x) |
| 42 | x = self.fc_select(x) |
| 43 | B, C, H, W = x.shape |
| 44 | x = x.view(B, self.num_paths, C // self.num_paths, H, W) |
| 45 | x = torch.softmax(x, dim=1) |
| 46 | return x |
| 47 | |
| 48 | |
| 49 | class SelectiveKernel(nn.Module): |