| 169 | self.gamma = nn.Parameter(torch.zeros(1)) |
| 170 | |
| 171 | def forward(self, x): |
| 172 | B, C, H, W = x.size() |
| 173 | |
| 174 | query = self.query(x).view(B, -1, H * W).permute(0, 2, 1) |
| 175 | key = self.key(x).view(B, -1, H * W) |
| 176 | value = self.value(x).view(B, -1, H * W) |
| 177 | |
| 178 | attention = F.softmax(torch.bmm(query, key), dim=-1) |
| 179 | out = torch.bmm(value, attention.permute(0, 2, 1)) |
| 180 | out = out.view(B, C, H, W) |
| 181 | |
| 182 | return self.gamma * out + x |
| 183 | |
| 184 | class LocalAttention(nn.Module): |
| 185 | def __init__(self, embed_dim, window_size=7, num_heads=8): |