| 815 | # |
| 816 | |
| 817 | class LuongAttnDecoderRNN(nn.Module): |
| 818 | def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1): |
| 819 | super(LuongAttnDecoderRNN, self).__init__() |
| 820 | |
| 821 | # Keep for reference |
| 822 | self.attn_model = attn_model |
| 823 | self.hidden_size = hidden_size |
| 824 | self.output_size = output_size |
| 825 | self.n_layers = n_layers |
| 826 | self.dropout = dropout |
| 827 | |
| 828 | # Define layers |
| 829 | self.embedding = embedding |
| 830 | self.embedding_dropout = nn.Dropout(dropout) |
| 831 | self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout)) |
| 832 | self.concat = nn.Linear(hidden_size * 2, hidden_size) |
| 833 | self.out = nn.Linear(hidden_size, output_size) |
| 834 | |
| 835 | self.attn = Attn(attn_model, hidden_size) |
| 836 | |
| 837 | def forward(self, input_step, last_hidden, encoder_outputs): |
| 838 | # Note: we run this one step (word) at a time |
| 839 | # Get embedding of current input word |
| 840 | embedded = self.embedding(input_step) |
| 841 | embedded = self.embedding_dropout(embedded) |
| 842 | # Forward through unidirectional GRU |
| 843 | rnn_output, hidden = self.gru(embedded, last_hidden) |
| 844 | # Calculate attention weights from the current GRU output |
| 845 | attn_weights = self.attn(rnn_output, encoder_outputs) |
| 846 | # Multiply attention weights to encoder outputs to get new "weighted sum" context vector |
| 847 | context = attn_weights.bmm(encoder_outputs.transpose(0, 1)) |
| 848 | # Concatenate weighted context vector and GRU output using Luong eq. 5 |
| 849 | rnn_output = rnn_output.squeeze(0) |
| 850 | context = context.squeeze(1) |
| 851 | concat_input = torch.cat((rnn_output, context), 1) |
| 852 | concat_output = torch.tanh(self.concat(concat_input)) |
| 853 | # Predict next word using Luong eq. 6 |
| 854 | output = self.out(concat_output) |
| 855 | output = F.softmax(output, dim=1) |
| 856 | # Return output and final hidden state |
| 857 | return output, hidden |
| 858 | |
| 859 | |
| 860 | ###################################################################### |