| 254 | |
| 255 | |
| 256 | class CGAttention(nn.Module): |
| 257 | def __init__(self, token_dim, vis_dim, hidden_dim, drop=0., bias=True) -> None: |
| 258 | super().__init__() |
| 259 | self.norm_v = nn.LayerNorm(vis_dim) |
| 260 | self.norm_t = nn.LayerNorm(token_dim) |
| 261 | self.q_proj = nn.Linear(token_dim, hidden_dim, bias=bias) |
| 262 | self.k_proj = nn.Linear(vis_dim, hidden_dim, bias=bias) |
| 263 | self.v_proj = nn.Linear(vis_dim, hidden_dim, bias=bias) |
| 264 | self.proj = nn.Linear(hidden_dim, token_dim) |
| 265 | self.proj_drop = nn.Dropout(drop) |
| 266 | self.norm = nn.LayerNorm(token_dim) |
| 267 | self.mlp = Mlp(token_dim, token_dim*2, token_dim, drop=drop) |
| 268 | self.tau = nn.Parameter(torch.ones(1), requires_grad=True) |
| 269 | |
| 270 | def with_pe(self, vis, pe): |
| 271 | return vis + pe |
| 272 | |
| 273 | def forward(self, tokens, vis, pe=None): |
| 274 | b, c, h , w = vis.shape |
| 275 | vis = rearrange(vis, 'b c h w -> b (h w) c') |
| 276 | if pe is not None: |
| 277 | vis = self.with_pe(vis, pe) |
| 278 | vis = self.norm_v(vis) |
| 279 | q = self.q_proj(self.norm_t(tokens)) |
| 280 | k = self.k_proj(vis) |
| 281 | v = self.v_proj(vis) |
| 282 | |
| 283 | q = l2norm(q, dim=-1) |
| 284 | k = l2norm(k, dim=-1) |
| 285 | raw_attn = (q @ k.transpose(-2, -1)) |
| 286 | tau = torch.clamp(self.tau, max=0).exp() |
| 287 | attn = gumbel_softmax(raw_attn, dim=-2, tau=tau) |
| 288 | hit_map = attn |
| 289 | attn = attn / (attn.sum(dim=-1, keepdim=True) + 1) |
| 290 | new_tokens = attn @ v |
| 291 | new_tokens = self.proj_drop(self.proj(new_tokens)) |
| 292 | new_tokens = self.mlp(self.norm(new_tokens+tokens)) |
| 293 | return new_tokens, hit_map.reshape(b, -1, h, w) |
| 294 | |
| 295 | class Decoder(nn.Module): |
| 296 | def __init__(self, args) -> None: |