-
Notifications
You must be signed in to change notification settings - Fork 0
/
attention_model.py
88 lines (78 loc) · 3.67 KB
/
attention_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.functional import dropout
from torch.autograd import Variable
import math
import random
class Attention(nn.Module):
def __init__(self, hidden_size):
super(Attention, self).__init__()
self.hidden_size = hidden_size
self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
self.v = nn.Parameter(torch.rand(hidden_size))
stdv = 1. / math.sqrt(self.v.size(0))
self.v.data.uniform_(-stdv, stdv)
def forward(self, hidden, encoder_outputs):
timestep = encoder_outputs.size(0)
h = hidden.repeat(timestep, 1, 1).transpose(0, 1)
encoder_outputs = encoder_outputs.transpose(0, 1) # [B*T*H]
attn_energies = self.score(h, encoder_outputs)
return F.softmax(attn_energies, dim=1).unsqueeze(1)
def score(self, hidden, encoder_outputs):
# [B*T*2H]->[B*T*H]
energy = F.relu(self.attn(torch.cat([hidden, encoder_outputs], 2)))
energy = energy.transpose(1, 2) # [B*H*T]
v = self.v.repeat(encoder_outputs.size(0), 1).unsqueeze(1) # [B*1*H]
energy = torch.bmm(v, energy) # [B*1*T]
return energy.squeeze(1) # [B*T]
class attn_Decoder(nn.Module):
def __init__(self, embed_size, hidden_size, output_size, dropout=0.2):
super(attn_Decoder, self).__init__()
self.embed_size = embed_size
self.hidden_size = hidden_size
self.output_size = output_size
self.embed = nn.Embedding(output_size, embed_size)
self.dropout = nn.Dropout(dropout, inplace=True)
self.attention = Attention(hidden_size)
self.gru = nn.GRU(hidden_size + embed_size, hidden_size,
dropout=dropout)
self.out = nn.Linear(hidden_size * 2, output_size)
def forward(self, input, last_hidden, encoder_outputs):
# Get the embedding of the current input word (last output word)
embedded = self.embed(input).unsqueeze(0) # (1,B,N)
embedded = self.dropout(embedded)
# Calculate attention weights and apply to encoder outputs
attn_weights = self.attention(last_hidden[-1], encoder_outputs)
context = attn_weights.bmm(encoder_outputs.transpose(0, 1)) # (B,1,N)
context = context.transpose(0, 1) # (1,B,N)
# Combine embedded input word and attended context, run through RNN
rnn_input = torch.cat([embedded, context], 2)
output, hidden = self.gru(rnn_input, last_hidden)
output = output.squeeze(0) # (1,B,N) -> (B,N)
context = context.squeeze(0)
output = self.out(torch.cat([output, context], 1))
output = F.log_softmax(output, dim=1)
return output, hidden, attn_weights
class attn_Seq2Seq(nn.Module):
def __init__(self, encoder, decoder,device):
super(attn_Seq2Seq, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.device=device
def forward(self, src, trg, teacher_forcing_ratio=0.5):
batch_size = src.size(1)
max_len = trg.size(0)
vocab_size = self.decoder.output_size
outputs = Variable(torch.zeros(max_len, batch_size, vocab_size)).to(self.device)
encoder_output, hidden,cell = self.encoder(src)
hidden = hidden[:1]
output = Variable(trg.data[0, :]) # sos
for t in range(1, max_len):
output, hidden, attn_weights = self.decoder(
output, hidden, encoder_output)
outputs[t] = output
is_teacher = random.random() < teacher_forcing_ratio
top1 = output.data.max(1)[1]
output = Variable(trg.data[t] if is_teacher else top1).cuda()
return outputs