| 741 | |
| 742 | # Luong attention layer |
| 743 | class Attn(nn.Module): |
| 744 | def __init__(self, method, hidden_size): |
| 745 | super(Attn, self).__init__() |
| 746 | self.method = method |
| 747 | if self.method not in ['dot', 'general', 'concat']: |
| 748 | raise ValueError(self.method, "is not an appropriate attention method.") |
| 749 | self.hidden_size = hidden_size |
| 750 | if self.method == 'general': |
| 751 | self.attn = nn.Linear(self.hidden_size, hidden_size) |
| 752 | elif self.method == 'concat': |
| 753 | self.attn = nn.Linear(self.hidden_size * 2, hidden_size) |
| 754 | self.v = nn.Parameter(torch.FloatTensor(hidden_size)) |
| 755 | |
| 756 | def dot_score(self, hidden, encoder_output): |
| 757 | return torch.sum(hidden * encoder_output, dim=2) |
| 758 | |
| 759 | def general_score(self, hidden, encoder_output): |
| 760 | energy = self.attn(encoder_output) |
| 761 | return torch.sum(hidden * energy, dim=2) |
| 762 | |
| 763 | def concat_score(self, hidden, encoder_output): |
| 764 | energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh() |
| 765 | return torch.sum(self.v * energy, dim=2) |
| 766 | |
| 767 | def forward(self, hidden, encoder_outputs): |
| 768 | # Calculate the attention weights (energies) based on the given method |
| 769 | if self.method == 'general': |
| 770 | attn_energies = self.general_score(hidden, encoder_outputs) |
| 771 | elif self.method == 'concat': |
| 772 | attn_energies = self.concat_score(hidden, encoder_outputs) |
| 773 | elif self.method == 'dot': |
| 774 | attn_energies = self.dot_score(hidden, encoder_outputs) |
| 775 | |
| 776 | # Transpose max_length and batch_size dimensions |
| 777 | attn_energies = attn_energies.t() |
| 778 | |
| 779 | # Return the softmax normalized probability scores (with added dimension) |
| 780 | return F.softmax(attn_energies, dim=1).unsqueeze(1) |
| 781 | |
| 782 | |
| 783 | ###################################################################### |