| 31 | |
| 32 | |
| 33 | class SpatialWeights(nn.Module): |
| 34 | def __init__(self, dim, reduction=1): |
| 35 | super(SpatialWeights, self).__init__() |
| 36 | self.dim = dim |
| 37 | self.mlp = nn.Sequential( |
| 38 | nn.Conv2d(self.dim * 2, self.dim // reduction, kernel_size=1), |
| 39 | nn.ReLU(inplace=True), |
| 40 | nn.Conv2d(self.dim // reduction, 2, kernel_size=1), |
| 41 | nn.Sigmoid(), |
| 42 | ) |
| 43 | |
| 44 | def forward(self, x1, x2): |
| 45 | B, _, H, W = x1.shape |
| 46 | x = torch.cat((x1, x2), dim=1) # B 2C H W |
| 47 | spatial_weights = self.mlp(x).reshape(B, 2, 1, H, W).permute(1, 0, 2, 3, 4) # 2 B 1 H W |
| 48 | return spatial_weights |
| 49 | |
| 50 | |
| 51 | # FRM |