-
Notifications
You must be signed in to change notification settings - Fork 4
/
seq2seq.py
385 lines (341 loc) · 18.6 KB
/
seq2seq.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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from config import *
from data_util import batch2TrainData, normalizeString, indexesFromSentence
class EncoderRNN(nn.Module):
def __init__(self, hidden_size, embedding, n_layers=1, dropout=0):
super(EncoderRNN, self).__init__()
self.n_layers = n_layers
self.hidden_size = hidden_size
self.embedding = embedding
# 初始化GRU,这里输入和hidden大小都是hidden_size,因为我们这里假设embedding层的输出大小是hidden_size
# 如果只有一层,那么不进行Dropout,否则使用传入的参数dropout进行GRU的Dropout。
self.gru = nn.GRU(hidden_size, hidden_size, n_layers,
dropout=(0 if n_layers == 1 else dropout), bidirectional=True)
def forward(self, input_seq, input_lengths, hidden=None):
# 输入是(max_length, batch),Embedding之后变成(max_length, batch, hidden_size)
embedded = self.embedding(input_seq)
# Pack padded batch of sequences for RNN module
# 因为RNN(GRU)需要知道实际的长度,所以PyTorch提供了一个函数pack_padded_sequence把输入向量和长度pack
# 到一个对象PackedSequence里,这样便于使用。
packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
# 通过GRU进行forward计算,需要传入输入和隐变量
# 如果传入的输入是一个Tensor (max_length, batch, hidden_size)
# 那么输出outputs是(max_length, batch, hidden_size*num_directions)。
# 第三维是hidden_size和num_directions的混合,它们实际排列顺序是num_directions在前面,因此我们可以
# 使用outputs.view(seq_len, batch, num_directions, hidden_size)得到4维的向量。
# 其中第三维是方向,第四位是隐状态。
# 而如果输入是PackedSequence对象,那么输出outputs也是一个PackedSequence对象,我们需要用
# 函数pad_packed_sequence把它变成一个shape为(max_length, batch, hidden*num_directions)的向量以及
# 一个list,表示输出的长度,当然这个list和输入的input_lengths完全一样,因此通常我们不需要它。
outputs, hidden = self.gru(packed, hidden)
# 参考前面的注释,我们得到outputs为(max_length, batch, hidden*num_directions)
outputs, _ = torch.nn.utils.rnn.pad_packed_sequence(outputs)
# 我们需要把输出的num_directions双向的向量加起来
# 因为outputs的第三维是先放前向的hidden_size个结果,然后再放后向的hidden_size个结果
# 所以outputs[:, :, :self.hidden_size]得到前向的结果
# outputs[:, :, self.hidden_size:]是后向的结果
# 注意,如果bidirectional是False,则outputs第三维的大小就是hidden_size,
# 这时outputs[:, : ,self.hidden_size:]是不存在的,因此也不会加上去。
# 对Python slicing不熟的读者可以看看下面的例子:
# >>> a=[1,2,3]
# >>> a[:3]
# [1, 2, 3]
# >>> a[3:]
# []
# >>> a[:3]+a[3:]
# [1, 2, 3]
# 这样就不用写下面的代码了:
# if bidirectional:
# outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:]
outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]
# 返回最终的输出和最后时刻的隐状态。
return outputs, hidden
# Luong 注意力layer
class Attn(torch.nn.Module):
def __init__(self, method, hidden_size):
super(Attn, self).__init__()
self.method = method
if self.method not in ['dot', 'general', 'concat']:
raise ValueError(self.method, "is not an appropriate attention method.")
self.hidden_size = hidden_size
if self.method == 'general':
self.attn = torch.nn.Linear(self.hidden_size, hidden_size)
elif self.method == 'concat':
self.attn = torch.nn.Linear(self.hidden_size * 2, hidden_size)
self.v = torch.nn.Parameter(torch.FloatTensor(hidden_size))
def dot_score(self, hidden, encoder_output):
# 输入hidden的shape是(1, batch=64, hidden_size=500)
# encoder_outputs的shape是(input_lengths=10, batch=64, hidden_size=500)
# hidden * encoder_output得到的shape是(10, 64, 500),然后对第3维求和就可以计算出score。
# dim定义要减少的维度
return torch.sum(hidden * encoder_output, dim=2)
def general_score(self, hidden, encoder_output):
energy = self.attn(encoder_output)
return torch.sum(hidden * energy, dim=2)
def concat_score(self, hidden, encoder_output):
energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()
return torch.sum(self.v * energy, dim=2)
# 输入是上一个时刻的隐状态hidden和所有时刻的Encoder的输出encoder_outputs
# 输出是注意力的概率,也就是长度为input_lengths的向量,它的和加起来是1。
def forward(self, hidden, encoder_outputs):
# 计算注意力的score,输入hidden的shape是(1, batch=64, hidden_size=500),表示t时刻batch数据的隐状态
# encoder_outputs的shape是(input_lengths=10, batch=64, hidden_size=500)
if self.method == 'general':
attn_energies = self.general_score(hidden, encoder_outputs)
elif self.method == 'concat':
attn_energies = self.concat_score(hidden, encoder_outputs)
elif self.method == 'dot':
# 计算内积,参考dot_score函数
attn_energies = self.dot_score(hidden, encoder_outputs)
# Transpose max_length and batch_size dimensions
# 把attn_energies从(max_length=10, batch=64)转置成(64, 10)
attn_energies = attn_energies.t()
# 使用softmax函数把score变成概率,shape仍然是(64, 10),然后用unsqueeze(1)变成
# (64, 1, 10),给指定位置加上维数为一的维度
return F.softmax(attn_energies, dim=1).unsqueeze(1)
class LuongAttnDecoderRNN(nn.Module):
def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):
super(LuongAttnDecoderRNN, self).__init__()
# Keep for reference
# 保存到self里,attn_model就是前面定义的Attn类的对象。
self.attn_model = attn_model
self.hidden_size = hidden_size
self.output_size = output_size
self.n_layers = n_layers
self.dropout = dropout
# Define layers
# 定义Decoder的layers
self.embedding = embedding
self.embedding_dropout = nn.Dropout(dropout)
self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout))
self.concat = nn.Linear(hidden_size * 2, hidden_size)
self.out = nn.Linear(hidden_size, output_size)
self.attn = Attn(attn_model, hidden_size)
def forward(self, input_step, last_hidden, encoder_outputs):
# 注意:decoder每一步只能处理一个时刻的数据,因为t时刻计算完了才能计算t+1时刻。
# input_step的shape是(1, 64),64是batch,1是当前输入的词ID(来自上一个时刻的输出)
# 通过embedding层变成(1, 64, 500),然后进行dropout,shape不变。
embedded = self.embedding(input_step)
embedded = self.embedding_dropout(embedded)
# 把embedded传入GRU进行forward计算
# 得到rnn_output的shape是(1, 64, 500)
# hidden是(2, 64, 500),因为是双向的GRU,所以第一维是2。
rnn_output, hidden = self.gru(embedded, last_hidden)
# 计算注意力权重, 根据前面的分析,attn_weights的shape是(64, 1, 10)
attn_weights = self.attn(rnn_output, encoder_outputs)
# encoder_outputs是(10, 64, 500)
# encoder_outputs.transpose(0, 1)后的shape是(64, 10, 500)
# attn_weights.bmm后是(64, 1, 500)
# bmm是批量的矩阵乘法,第一维是batch,我们可以把attn_weights看成64个(1,10)的矩阵
# 把encoder_outputs.transpose(0, 1)看成64个(10, 500)的矩阵
# 那么bmm就是64个(1, 10)矩阵 x (10, 500)矩阵,最终得到(64, 1, 500)
context = attn_weights.bmm(encoder_outputs.transpose(0, 1))
# 把context向量和GRU的输出拼接起来
# rnn_output从(1, 64, 500)变成(64, 500)
rnn_output = rnn_output.squeeze(0)
# context从(64, 1, 500)变成(64, 500)
context = context.squeeze(1)
# 拼接得到(64, 1000)
concat_input = torch.cat((rnn_output, context), 1)
# self.concat是一个矩阵(1000, 500),
# self.concat(concat_input)的输出是(64, 500)
# 然后用tanh把输出返回变成(-1,1),concat_output的shape是(64, 500)
concat_output = torch.tanh(self.concat(concat_input))
# out是(500, 词典大小=7826)
output = self.out(concat_output)
# 用softmax变成概率,表示当前时刻输出每个词的概率。
output = F.softmax(output, dim=1)
# 返回 output和新的隐状态
return output, hidden
# Masked loss
# ~~~~~~~~~~~
# Decoder输出的是每个词的概率分布,因此可以使用交叉熵损失函数。
# 因为是一个batch的数据里有一些是padding的,因此这些位置的预测是没有必要计算loss的,
# 因此我们需要使用前面的mask矩阵把对应位置的loss去掉,我们可以通过下面的函数来实现计算Masked的loss。
def maskNLLLoss(inp, target, mask):
# 计算实际的词的个数,因为padding是0,非padding是1,因此sum就可以得到词的个数
nTotal = mask.sum()
# torch.gather函数首先把0.4和0.3(正确分类对应的概率值)选出来
crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)).squeeze(1))
loss = crossEntropy.masked_select(mask).mean()
loss = loss.to(device)
return loss, nTotal.item()
# PyTorch的RNN模块(RNN, LSTM, GRU)也可以当成普通的非循环的网络来使用。
# 在Encoder部分,我们是直接把所有时刻的数据都传入RNN,让它一次计算出所有的结果,
# 但是在Decoder的时候(非teacher forcing)后一个时刻的输入来自前一个时刻的输出,因此无法一次计算
def train(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder, embedding,
encoder_optimizer, decoder_optimizer, batch_size, clip, max_length=MAX_LENGTH):
# 梯度清空
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
# 设置device,从而支持GPU,当然如果没有GPU也能工作。
input_variable = input_variable.to(device)
lengths = lengths.to(device)
target_variable = target_variable.to(device)
mask = mask.to(device)
# 初始化变量
loss = 0
print_losses = []
n_totals = 0
# encoder的Forward计算
encoder_outputs, encoder_hidden = encoder(input_variable, lengths)
# Decoder的初始输入是SOS,我们需要构造(1, batch)的输入,表示第一个时刻batch个输入。
decoder_input = torch.LongTensor([[SOS_token for _ in range(batch_size)]])
decoder_input = decoder_input.to(device)
# 注意:Encoder是双向的,而Decoder是单向的,因此从下往上取n_layers个
decoder_hidden = encoder_hidden[:decoder.n_layers]
# 确定是否teacher forcing
use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False
# 一次处理一个时刻
if use_teacher_forcing:
for t in range(max_target_len):
decoder_output, decoder_hidden = decoder(
decoder_input, decoder_hidden, encoder_outputs
)
# Teacher forcing: 下一个时刻的输入是当前正确答案
decoder_input = target_variable[t].view(1, -1)
# 计算累计的loss
mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])
loss += mask_loss
print_losses.append(mask_loss.item() * nTotal)
n_totals += nTotal
else:
for t in range(max_target_len):
decoder_output, decoder_hidden = decoder(
decoder_input, decoder_hidden, encoder_outputs
)
# 不是teacher forcing: 下一个时刻的输入是当前模型预测概率最高的值
_, topi = decoder_output.topk(1)
decoder_input = torch.LongTensor([[topi[i][0] for i in range(batch_size)]])
decoder_input = decoder_input.to(device)
# 计算累计的loss
mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])
loss += mask_loss
print_losses.append(mask_loss.item() * nTotal)
n_totals += nTotal
# 反向计算
loss.backward()
# 对encoder和decoder进行梯度裁剪
_ = torch.nn.utils.clip_grad_norm_(encoder.parameters(), clip)
_ = torch.nn.utils.clip_grad_norm_(decoder.parameters(), clip)
# 更新参数
encoder_optimizer.step()
decoder_optimizer.step()
return sum(print_losses) / n_totals
# def trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer, decoder_optimizer, embedding, encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size, print_every, save_every, clip, corpus_name, loadFilename):
#
# # 随机选择n_iteration个batch的数据(pair)
# training_batches = [batch2TrainData(voc, [random.choice(pairs) for _ in range(batch_size)])
# for _ in range(n_iteration)]
#
# # 初始化
# print('Initializing ...')
# start_iteration = 1
# print_loss = 0
# if loadFilename:
# start_iteration = checkpoint['iteration'] + 1
#
# # 训练
# print("Training...")
# for iteration in range(start_iteration, n_iteration + 1):
# training_batch = training_batches[iteration - 1]
# # Extract fields from batch
# input_variable, lengths, target_variable, mask, max_target_len = training_batch
#
# # 训练一个batch的数据
# loss = train(input_variable, lengths, target_variable, mask, max_target_len, encoder,
# decoder, embedding, encoder_optimizer, decoder_optimizer, batch_size, clip)
# print_loss += loss
#
# # 进度
# if iteration % print_every == 0:
# print_loss_avg = print_loss / print_every
# print("Iteration: {}; Percent complete: {:.1f}%; Average loss: {:.4f}".format(iteration, iteration / n_iteration * 100, print_loss_avg))
# print_loss = 0
#
# # 保存checkpoint
# if (iteration % save_every == 0):
# if not os.path.exists(directory):
# os.makedirs(directory)
# torch.save({
# 'iteration': iteration,
# 'en': encoder.state_dict(),
# 'de': decoder.state_dict(),
# 'en_opt': encoder_optimizer.state_dict(),
# 'de_opt': decoder_optimizer.state_dict(),
# 'loss': loss,
# 'voc_dict': voc.__dict__,
# 'embedding': embedding.state_dict()
# }, os.path.join(directory, '{}_{}.tar'.format(iteration, 'checkpoint')))
# class GreedySearchDecoder(nn.Module):
# def __init__(self, encoder, decoder):
# super(GreedySearchDecoder, self).__init__()
# self.encoder = encoder
# self.decoder = decoder
#
# def forward(self, input_seq, input_length, max_length):
# # Encoder的Forward计算
# encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)
# # 把Encoder最后时刻的隐状态作为Decoder的初始值
# decoder_hidden = encoder_hidden[:decoder.n_layers]
# # 因为我们的函数都是要求(time,batch),因此即使只有一个数据,也要做出二维的。
# # Decoder的初始输入是SOS
# decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token
# # 用于保存解码结果的tensor
# all_tokens = torch.zeros([0], device=device, dtype=torch.long)
# all_scores = torch.zeros([0], device=device)
# # 循环,这里只使用长度限制,后面处理的时候把EOS去掉了。
# for _ in range(max_length):
# # Decoder forward一步
# decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
# # decoder_outputs是(batch=1, vob_size)
# # 使用max返回概率最大的词和得分
# decoder_scores, decoder_input = torch.max(decoder_output, dim=1)
# # 把解码结果保存到all_tokens和all_scores里
# all_tokens = torch.cat((all_tokens, decoder_input), dim=0)
# all_scores = torch.cat((all_scores, decoder_scores), dim=0)
# # decoder_input是当前时刻输出的词的ID,这是个一维的向量,因为max会减少一维。
# # 但是decoder要求有一个batch维度,因此用unsqueeze增加batch维度。
# decoder_input = torch.unsqueeze(decoder_input, 0)
# # 返回所有的词和得分。
# return all_tokens, all_scores
def evaluate(encoder, decoder, searcher, voc, sentence, max_length=MAX_LENGTH):
### 把输入的一个batch句子变成id
indexes_batch = [indexesFromSentence(voc, sentence)]
# 创建lengths tensor
lengths = torch.tensor([len(indexes) for indexes in indexes_batch])
# 转置
input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)
# 放到合适的设备上(比如GPU)
input_batch = input_batch.to(device)
lengths = lengths.to(device)
# 用searcher解码
tokens, scores = searcher(input_batch, lengths, max_length)
# ID变成词。
decoded_words = [voc.index2word[token.item()] for token in tokens]
return decoded_words
def evaluateInput(encoder, decoder, searcher, voc):
input_sentence = ''
while(1):
try:
# 得到用户终端的输入
input_sentence = input('> ')
# 是否退出
if input_sentence == 'q' or input_sentence == 'quit': break
# 句子归一化
input_sentence = normalizeString(input_sentence)
# 生成响应Evaluate sentence
output_words = evaluate(encoder, decoder, searcher, voc, input_sentence)
# 去掉EOS后面的内容
words = []
for word in output_words:
if word == 'EOS':
break
elif word != 'PAD':
words.append(word)
print('Bot:', ' '.join(words))
except KeyError:
print("Error: Encountered unknown word.")