(self, tokens, vis, pe=None)
| 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: |
nothing calls this directly
no test coverage detected