| 10 | |
| 11 | |
| 12 | class DecoderLayer(nn.Module): |
| 13 | def __init__(self, dim, num_heads, mlp_ratio, dropout) -> None: |
| 14 | super().__init__() |
| 15 | self.dim = dim |
| 16 | |
| 17 | self.r2r_attn = nn.MultiheadAttention( |
| 18 | dim, num_heads, dropout=dropout, batch_first=True |
| 19 | ) |
| 20 | self.m2m_attn = nn.MultiheadAttention( |
| 21 | dim, num_heads, dropout=dropout, batch_first=True |
| 22 | ) |
| 23 | self.cross_attn = nn.MultiheadAttention( |
| 24 | dim, num_heads, dropout=dropout, batch_first=True |
| 25 | ) |
| 26 | |
| 27 | self.ffn = nn.Sequential( |
| 28 | nn.Linear(dim, dim * mlp_ratio), |
| 29 | nn.ReLU(inplace=True), |
| 30 | nn.Dropout(dropout), |
| 31 | nn.Linear(dim * mlp_ratio, dim), |
| 32 | ) |
| 33 | |
| 34 | self.norm1 = nn.LayerNorm(dim) |
| 35 | self.norm2 = nn.LayerNorm(dim) |
| 36 | self.norm3 = nn.LayerNorm(dim) |
| 37 | self.norm4 = nn.LayerNorm(dim) |
| 38 | self.dropout1 = nn.Dropout(dropout) |
| 39 | self.dropout2 = nn.Dropout(dropout) |
| 40 | self.dropout3 = nn.Dropout(dropout) |
| 41 | |
| 42 | def forward( |
| 43 | self, |
| 44 | tgt, |
| 45 | memory, |
| 46 | tgt_key_padding_mask: Optional[Tensor] = None, |
| 47 | memory_key_padding_mask: Optional[Tensor] = None, |
| 48 | m_pos: Optional[Tensor] = None, |
| 49 | ): |
| 50 | """ |
| 51 | tgt: (bs, R, M, dim) |
| 52 | tgt_key_padding_mask: (bs, R) |
| 53 | """ |
| 54 | bs, R, M, D = tgt.shape |
| 55 | |
| 56 | tgt = tgt.transpose(1, 2).reshape(bs * M, R, D) |
| 57 | tgt2 = self.norm1(tgt) |
| 58 | tgt2 = self.r2r_attn( |
| 59 | tgt2, tgt2, tgt2, key_padding_mask=tgt_key_padding_mask.repeat(M, 1) |
| 60 | )[0] |
| 61 | tgt = tgt + self.dropout1(tgt2) |
| 62 | |
| 63 | tgt_tmp = tgt.reshape(bs, M, R, D).transpose(1, 2).reshape(bs * R, M, D) |
| 64 | tgt_valid_mask = ~tgt_key_padding_mask.reshape(-1) |
| 65 | tgt_valid = tgt_tmp[tgt_valid_mask] |
| 66 | tgt2_valid = self.norm2(tgt_valid) |
| 67 | tgt2_valid, _ = self.m2m_attn( |
| 68 | tgt2_valid + m_pos, tgt2_valid + m_pos, tgt2_valid |
| 69 | ) |