| 171 | |
| 172 | |
| 173 | class DecoderLayer(nn.Module): |
| 174 | def __init__(self, d_model, n_head, dropout): |
| 175 | super().__init__() |
| 176 | self.self_attn_norm = nn.LayerNorm(d_model) |
| 177 | self.self_attn = DecoderMultiHeadAttention(d_model, n_head, dropout) |
| 178 | |
| 179 | self.cross_attn_norm = nn.LayerNorm(d_model) |
| 180 | self.cross_attn = DecoderMultiHeadAttention(d_model, n_head, dropout) |
| 181 | |
| 182 | self.mlp_norm = nn.LayerNorm(d_model) |
| 183 | self.mlp = PositionwiseFeedForward(d_model, d_model*4, dropout) |
| 184 | |
| 185 | def forward(self, dec_input, enc_output, self_attn_mask, cross_attn_mask, |
| 186 | cache=None): |
| 187 | x = dec_input |
| 188 | residual = x |
| 189 | x = self.self_attn_norm(x) |
| 190 | if cache is not None: |
| 191 | xq = x[:, -1:, :] |
| 192 | residual = residual[:, -1:, :] |
| 193 | self_attn_mask = self_attn_mask[:, -1:, :] |
| 194 | else: |
| 195 | xq = x |
| 196 | x = self.self_attn(xq, x, x, mask=self_attn_mask) |
| 197 | x = residual + x |
| 198 | |
| 199 | residual = x |
| 200 | x = self.cross_attn_norm(x) |
| 201 | x = self.cross_attn(x, enc_output, enc_output, mask=cross_attn_mask) |
| 202 | x = residual + x |
| 203 | |
| 204 | residual = x |
| 205 | x = self.mlp_norm(x) |
| 206 | x = residual + self.mlp(x) |
| 207 | |
| 208 | if cache is not None: |
| 209 | x = torch.cat([cache, x], dim=1) |
| 210 | |
| 211 | return x |
| 212 | |
| 213 | |
| 214 | class DecoderMultiHeadAttention(nn.Module): |