Forward pass for a single stream (i.e. content or query) tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency. Both tgt_kv and memory are expected to be LayerNorm'd too. memory is LayerNorm'd by ViT.
(
self,
tgt: Tensor,
tgt_norm: Tensor,
tgt_kv: Tensor,
memory: Tensor,
tgt_mask: Optional[Tensor],
tgt_key_padding_mask: Optional[Tensor],
)
| 69 | super().__setstate__(state) |
| 70 | |
| 71 | def forward_stream( |
| 72 | self, |
| 73 | tgt: Tensor, |
| 74 | tgt_norm: Tensor, |
| 75 | tgt_kv: Tensor, |
| 76 | memory: Tensor, |
| 77 | tgt_mask: Optional[Tensor], |
| 78 | tgt_key_padding_mask: Optional[Tensor], |
| 79 | ): |
| 80 | """Forward pass for a single stream (i.e. content or query) tgt_norm is |
| 81 | just a LayerNorm'd tgt. |
| 82 | |
| 83 | Added as a separate parameter for efficiency. Both tgt_kv and memory |
| 84 | are expected to be LayerNorm'd too. memory is LayerNorm'd by ViT. |
| 85 | """ |
| 86 | tgt2, sa_weights = self.self_attn( |
| 87 | tgt_norm, |
| 88 | tgt_kv, |
| 89 | tgt_kv, |
| 90 | attn_mask=tgt_mask, |
| 91 | key_padding_mask=tgt_key_padding_mask) |
| 92 | tgt = tgt + self.dropout1(tgt2) |
| 93 | |
| 94 | tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory) |
| 95 | self.attn_map = ca_weights |
| 96 | tgt = tgt + self.dropout2(tgt2) |
| 97 | |
| 98 | tgt2 = self.linear2( |
| 99 | self.dropout(self.activation(self.linear1(self.norm2(tgt))))) |
| 100 | tgt = tgt + self.dropout3(tgt2) |
| 101 | return tgt, sa_weights, ca_weights |
| 102 | |
| 103 | def forward( |
| 104 | self, |