r"""TransformerEncoderLayer is made up of self-attn and feedforward network. This standard encoder layer is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017
| 16 | |
| 17 | |
| 18 | class TransformerEncoderLayer(Module): |
| 19 | r"""TransformerEncoderLayer is made up of self-attn and feedforward network. |
| 20 | This standard encoder layer is based on the paper "Attention Is All You Need". |
| 21 | Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, |
| 22 | Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in |
| 23 | Neural Information Processing Systems, pages 6000-6010. Users may modify or implement |
| 24 | in a different way during application. |
| 25 | Args: |
| 26 | d_model: the number of expected features in the input (required). |
| 27 | nhead: the number of heads in the multiheadattention models (required). |
| 28 | dim_feedforward: the dimension of the feedforward network model (default=2048). |
| 29 | dropout: the dropout value (default=0.1). |
| 30 | activation: the activation function of intermediate layer, relu or gelu (default=relu). |
| 31 | Examples:: |
| 32 | >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8) |
| 33 | >>> src = torch.rand(10, 32, 512) |
| 34 | >>> out = encoder_layer(src) |
| 35 | """ |
| 36 | |
| 37 | def __init__(self, d_model, nhead, bidirectional=True, dropout=0, activation="relu"): |
| 38 | super(TransformerEncoderLayer, self).__init__() |
| 39 | |
| 40 | |
| 41 | #self.self_attn = PerformerSelfAttention(d_model, nhead) |
| 42 | # Implementation of Feedforward model |
| 43 | #dim_feedforward = d_model * 2 |
| 44 | #self.linear1 = Linear(d_model, dim_feedforward) |
| 45 | self.gru = GRU(d_model, d_model*2, 1, bidirectional=bidirectional) |
| 46 | self.dropout = Dropout(dropout) |
| 47 | #self.linear2 = Linear(dim_feedforward, d_model) |
| 48 | |
| 49 | if bidirectional: |
| 50 | self.linear2 = Linear(d_model*2*2, d_model) |
| 51 | self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout) |
| 52 | self.norm1 = InstantLayerNorm1d(d_model) |
| 53 | self.dropout1 = Dropout(dropout) |
| 54 | else: |
| 55 | self.linear2 = Linear(d_model*2, d_model) |
| 56 | |
| 57 | self.norm2 = InstantLayerNorm1d(d_model) |
| 58 | self.dropout2 = Dropout(dropout) |
| 59 | |
| 60 | self.activation = _get_activation_fn(activation) |
| 61 | self.bidirectional = bidirectional |
| 62 | |
| 63 | def __setstate__(self, state): |
| 64 | if 'activation' not in state: |
| 65 | state['activation'] = F.relu |
| 66 | super(TransformerEncoderLayer, self).__setstate__(state) |
| 67 | |
| 68 | def forward(self, src, src_mask=None, src_key_padding_mask=None): |
| 69 | # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor |
| 70 | r"""Pass the input through the encoder layer. |
| 71 | Args: |
| 72 | src: the sequnce to the encoder layer (required). |
| 73 | src_mask: the mask for the src sequence (optional). |
| 74 | src_key_padding_mask: the mask for the src keys per batch (optional). |
| 75 | Shape: |