| 84 | |
| 85 | |
| 86 | class CBAMBlock(nn.Module): |
| 87 | def __init__(self, channel=512,reduction=16,kernel_size=7, use_cls_token=False): |
| 88 | super().__init__() |
| 89 | self.ca = ChannelAttention(channel=channel,reduction=reduction) |
| 90 | self.sa = SpatialAttention(kernel_size=kernel_size) |
| 91 | self.use_cls_token = use_cls_token |
| 92 | |
| 93 | |
| 94 | def init_weights(self): |
| 95 | for m in self.modules(): |
| 96 | if isinstance(m, nn.Conv2d): |
| 97 | init.kaiming_normal_(m.weight, mode='fan_out') |
| 98 | if m.bias is not None: |
| 99 | init.constant_(m.bias, 0) |
| 100 | elif isinstance(m, nn.BatchNorm2d): |
| 101 | init.constant_(m.weight, 1) |
| 102 | init.constant_(m.bias, 0) |
| 103 | elif isinstance(m, nn.Linear): |
| 104 | init.normal_(m.weight, std=0.001) |
| 105 | if m.bias is not None: |
| 106 | init.constant_(m.bias, 0) |
| 107 | |
| 108 | def forward(self, x, H, W): |
| 109 | # import pdb; pdb.set_trace() |
| 110 | if self.use_cls_token: |
| 111 | cls_feat = x[:, :1, :] |
| 112 | x = x[:, 1:, :] |
| 113 | B = x.shape[0] |
| 114 | x = x.permute(0, 2, 1).reshape(B, -1, H, W) |
| 115 | b, c, _, _ = x.size() |
| 116 | residual = x |
| 117 | out = x * self.ca(x) |
| 118 | out = out * self.sa(out) + residual |
| 119 | if self.use_cls_token: |
| 120 | return torch.cat([cls_feat, out.reshape(B, -1, H * W).permute(0, 2, 1)], dim=1) |
| 121 | else: |
| 122 | return out.reshape(B, -1, H * W).permute(0, 2, 1) |
| 123 | |
| 124 | |
| 125 | class SEAttention(nn.Module): |