| 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 |