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