| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class SpatialAttention(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.conv = nn.Sequential( |
| nn.Conv2d(2, 1, kernel_size=(1, 1), stride=1), nn.BatchNorm2d(1), nn.ReLU() |
| ) |
|
|
| self.sgap = nn.AvgPool2d(2) |
|
|
| def forward(self, x): |
| B, H, W, C = x.shape |
| x = x.reshape(B, C, H, W) |
|
|
| mx = torch.max(x, 1)[0].unsqueeze(1) |
| avg = torch.mean(x, 1).unsqueeze(1) |
| combined = torch.cat([mx, avg], dim=1) |
| fmap = self.conv(combined) |
| weight_map = torch.sigmoid(fmap) |
| out = (x * weight_map).mean(dim=(-2, -1)) |
|
|
| return out, x * weight_map |
|
|
|
|
| class TokenLearner(nn.Module): |
| def __init__(self, S) -> None: |
| super().__init__() |
| self.S = S |
| self.tokenizers = nn.ModuleList([SpatialAttention() for _ in range(S)]) |
|
|
| def forward(self, x): |
| B, _, _, C = x.shape |
| Z = torch.Tensor(B, self.S, C).to(x) |
| for i in range(self.S): |
| Ai, _ = self.tokenizers[i](x) |
| Z[:, i, :] = Ai |
| return Z |
|
|
|
|
| class TokenFuser(nn.Module): |
| def __init__(self, H, W, C, S) -> None: |
| super().__init__() |
| self.projection = nn.Linear(S, S, bias=False) |
| self.Bi = nn.Linear(C, S) |
| self.spatial_attn = SpatialAttention() |
| self.S = S |
|
|
| def forward(self, y, x): |
| B, S, C = y.shape |
| B, H, W, C = x.shape |
|
|
| Y = self.projection(y.reshape(B, C, S)).reshape(B, S, C) |
| Bw = torch.sigmoid(self.Bi(x)).reshape(B, H * W, S) |
| BwY = torch.matmul(Bw, Y) |
|
|
| _, xj = self.spatial_attn(x) |
| xj = xj.reshape(B, H * W, C) |
|
|
| out = (BwY + xj).reshape(B, H, W, C) |
|
|
| return out |
|
|
