| import torch | |
| import torch.nn as nn | |
| class Attention(nn.Module): | |
| def __init__(self, hidden_size): | |
| super(Attention, self).__init__() | |
| self.W1 = nn.Linear(hidden_size, hidden_size) | |
| self.W2 = nn.Linear(hidden_size, hidden_size) | |
| self.v = nn.Linear(hidden_size, 1, bias=False) | |
| def forward(self, hidden, encoder_outputs): | |
| sequence_len = encoder_outputs.shape[1] | |
| hidden = hidden.unsqueeze(1).repeat(1, sequence_len, 1) | |
| energy = torch.tanh(self.W1(encoder_outputs) + self.W2(hidden)) | |
| attention = self.v(energy).squeeze(2) | |
| attention_weights = torch.softmax(attention, dim=1) | |
| context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1) | |
| return context, attention_weights | |
| class SimpleRecurrentNetworkWithAttention(nn.Module): | |
| def __init__(self, input_size, hidden_size, output_size, cell_type='RNN', device='cpu'): | |
| super(SimpleRecurrentNetworkWithAttention, self).__init__() | |
| self.device = device | |
| self.embedding = nn.Embedding(input_size, hidden_size) | |
| self.attention = Attention(hidden_size * 2) | |
| if cell_type == 'LSTM': | |
| self.rnn = nn.LSTM(hidden_size, hidden_size, batch_first=True, bidirectional=True) | |
| elif cell_type == 'GRU': | |
| self.rnn = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True) | |
| else: | |
| self.rnn = nn.RNN(hidden_size, hidden_size, batch_first=True, bidirectional=True) | |
| self.fc = nn.Linear(hidden_size * 2, output_size) | |
| def forward(self, inputs): | |
| embedded = self.embedding(inputs.to(self.device)) | |
| rnn_output, hidden = self.rnn(embedded) | |
| if isinstance(hidden, tuple): | |
| hidden = hidden[0] | |
| hidden = torch.cat((hidden[-2], hidden[-1]), dim=1) | |
| context, attention_weights = self.attention(hidden, rnn_output) | |
| output = self.fc(context) | |
| return output, attention_weights | |