| 229 | |
| 230 | |
| 231 | class TransformerDecoderLayer(nn.Module): |
| 232 | def __init__( |
| 233 | self, |
| 234 | d_model, |
| 235 | nhead, |
| 236 | dim_feedforward=2048, |
| 237 | dropout=0.1, |
| 238 | activation="relu", |
| 239 | normalize_before=False, |
| 240 | ): |
| 241 | super().__init__() |
| 242 | self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
| 243 | self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
| 244 | # Implementation of Feedforward model |
| 245 | self.linear1 = nn.Linear(d_model, dim_feedforward) |
| 246 | self.dropout = nn.Dropout(dropout) |
| 247 | self.linear2 = nn.Linear(dim_feedforward, d_model) |
| 248 | |
| 249 | self.norm1 = nn.LayerNorm(d_model) |
| 250 | self.norm2 = nn.LayerNorm(d_model) |
| 251 | self.norm3 = nn.LayerNorm(d_model) |
| 252 | self.dropout1 = nn.Dropout(dropout) |
| 253 | self.dropout2 = nn.Dropout(dropout) |
| 254 | self.dropout3 = nn.Dropout(dropout) |
| 255 | |
| 256 | self.activation = _get_activation_fn(activation) |
| 257 | self.normalize_before = normalize_before |
| 258 | |
| 259 | def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
| 260 | return tensor if pos is None else tensor + pos |
| 261 | |
| 262 | def forward_post( |
| 263 | self, |
| 264 | tgt, |
| 265 | memory, |
| 266 | tgt_mask: Optional[Tensor] = None, |
| 267 | memory_mask: Optional[Tensor] = None, |
| 268 | tgt_key_padding_mask: Optional[Tensor] = None, |
| 269 | memory_key_padding_mask: Optional[Tensor] = None, |
| 270 | pos: Optional[Tensor] = None, |
| 271 | query_pos: Optional[Tensor] = None, |
| 272 | ): |
| 273 | q = k = self.with_pos_embed(tgt, query_pos) |
| 274 | tgt2 = self.self_attn( |
| 275 | q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask |
| 276 | )[0] |
| 277 | tgt = tgt + self.dropout1(tgt2) |
| 278 | tgt = self.norm1(tgt) |
| 279 | tgt2 = self.multihead_attn( |
| 280 | query=self.with_pos_embed(tgt, query_pos), |
| 281 | key=self.with_pos_embed(memory, pos), |
| 282 | value=memory, |
| 283 | attn_mask=memory_mask, |
| 284 | key_padding_mask=memory_key_padding_mask, |
| 285 | )[0] |
| 286 | tgt = tgt + self.dropout2(tgt2) |
| 287 | tgt = self.norm2(tgt) |
| 288 | tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) |