A Transformer decoder layer supporting two-stream attention (XLNet) This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch.
| 26 | |
| 27 | |
| 28 | class DecoderLayer(nn.Module): |
| 29 | """A Transformer decoder layer supporting two-stream attention (XLNet) This |
| 30 | implements a pre-LN decoder, as opposed to the post-LN default in |
| 31 | PyTorch.""" |
| 32 | |
| 33 | def __init__( |
| 34 | self, |
| 35 | d_model, |
| 36 | nhead, |
| 37 | dim_feedforward=2048, |
| 38 | dropout=0.1, |
| 39 | activation='gelu', |
| 40 | layer_norm_eps=1e-5, |
| 41 | ): |
| 42 | super().__init__() |
| 43 | self.self_attn = nn.MultiheadAttention(d_model, |
| 44 | nhead, |
| 45 | dropout=dropout, |
| 46 | batch_first=True) |
| 47 | self.cross_attn = nn.MultiheadAttention(d_model, |
| 48 | nhead, |
| 49 | dropout=dropout, |
| 50 | batch_first=True) |
| 51 | # Implementation of Feedforward model |
| 52 | self.linear1 = nn.Linear(d_model, dim_feedforward) |
| 53 | self.dropout = nn.Dropout(dropout) |
| 54 | self.linear2 = nn.Linear(dim_feedforward, d_model) |
| 55 | |
| 56 | self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps) |
| 57 | self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps) |
| 58 | self.norm_q = nn.LayerNorm(d_model, eps=layer_norm_eps) |
| 59 | self.norm_c = nn.LayerNorm(d_model, eps=layer_norm_eps) |
| 60 | self.dropout1 = nn.Dropout(dropout) |
| 61 | self.dropout2 = nn.Dropout(dropout) |
| 62 | self.dropout3 = nn.Dropout(dropout) |
| 63 | |
| 64 | self.activation = transformer._get_activation_fn(activation) |
| 65 | |
| 66 | def __setstate__(self, state): |
| 67 | if 'activation' not in state: |
| 68 | state['activation'] = F.gelu |
| 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 | """ |