| 7 | |
| 8 | # Feature Rectify Module |
| 9 | class ChannelWeights(nn.Module): |
| 10 | def __init__(self, dim, reduction=1): |
| 11 | super(ChannelWeights, self).__init__() |
| 12 | self.dim = dim |
| 13 | self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| 14 | self.max_pool = nn.AdaptiveMaxPool2d(1) |
| 15 | self.mlp = nn.Sequential( |
| 16 | nn.Linear(self.dim * 4, self.dim * 4 // reduction), |
| 17 | nn.ReLU(inplace=True), |
| 18 | nn.Linear(self.dim * 4 // reduction, self.dim * 2), |
| 19 | nn.Sigmoid(), |
| 20 | ) |
| 21 | |
| 22 | def forward(self, x1, x2): |
| 23 | B, _, H, W = x1.shape |
| 24 | x = torch.cat((x1, x2), dim=1) |
| 25 | avg = self.avg_pool(x).view(B, self.dim * 2) |
| 26 | max = self.max_pool(x).view(B, self.dim * 2) |
| 27 | y = torch.cat((avg, max), dim=1) # B 4C |
| 28 | y = self.mlp(y).view(B, self.dim * 2, 1) |
| 29 | channel_weights = y.reshape(B, 2, self.dim, 1, 1).permute(1, 0, 2, 3, 4) # 2 B C 1 1 |
| 30 | return channel_weights |
| 31 | |
| 32 | |
| 33 | class SpatialWeights(nn.Module): |