| 187 | return out |
| 188 | |
| 189 | class OriLoadToken(nn.Module): |
| 190 | def __init__(self, token_dim, bias, drop) -> None: |
| 191 | super().__init__() |
| 192 | self.cross_attn = CrossAttn( |
| 193 | q_dim=token_dim, |
| 194 | kv_dim=768, |
| 195 | hidden_dim=token_dim, |
| 196 | num_heads=1, |
| 197 | out_dim=token_dim, |
| 198 | qkv_bias=bias, |
| 199 | attn_drop=drop, |
| 200 | proj_drop=drop, |
| 201 | ) |
| 202 | self.normq = nn.LayerNorm(token_dim) |
| 203 | self.normk = nn.LayerNorm(768) |
| 204 | |
| 205 | self.normq = nn.LayerNorm(token_dim) |
| 206 | self.normk = nn.LayerNorm(768) |
| 207 | |
| 208 | def forward(self, tokens, text, pad_mask): |
| 209 | tokens = tokens + self.cross_attn(query=self.normq(tokens), key=self.normk(text.permute(0,2,1)), mask=pad_mask[...,0]) |
| 210 | return tokens |
| 211 | |
| 212 | # updated version |
| 213 | class LoadToken(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected