使用注意力机制的Seq2Seq

引入

在使用注意力机制的Seq2Seq模型中，Key和Value都是编码器最后一层所有时间步的隐藏单元数据，编码器使用双向循环神经网络，num_layers=l，batch_size=n，hidden_size=h，src_len=s，则： \[ \mathbf{K} = \mathbf{V} = \mathbf{H}^{(l)}\in\mathbb{R}^{n\times s\times 2h} \] Query是解码器当前时间步\(t\)的输入隐藏层数据\(\mathbf{H}_{t-1}\in\mathbb{R}^{n\times h}\)

编码器将最后一个时间步的最后一个隐藏单元数据传递给解码器充当隐藏层数据，解码器的输入为target_embedded（如果是训练则为真实target，如果为预测则为上一个时间步的输出预测）和注意力attention，全连接层预测的输入为target_embedded+attention+hidden（隐藏层输出）

结构

编码器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 编码器
class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, enc_hidden_size, num_layers=1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, enc_hidden_size, num_layers=num_layers, bidirectional=True, batch_first=True)
        self.fc = nn.Linear(enc_hidden_size*2, enc_hidden_size)
    
    def forward(self, src, src_len):
        embedded = self.embedding(src)
        packed = pack_padded_sequence(embedded, src_len.cpu(), batch_first=True, enforce_sorted=True)
        outputs, hidden = self.rnn(packed)
        outputs, _ = pad_packed_sequence(outputs, batch_first=True)
        # outputs: [batch_size, src_len, enc_hidden_size*2]
        
        # 合并双向最后层的隐藏状态
        hidden = torch.tanh(self.fc(torch.cat((hidden[-1], hidden[-2]), dim=1)))
        # hidden: [batch_size, enc_hidden_size]
        return outputs, hidden

注意力机制

使用加性注意力机制 \[ \begin{aligned} &f(\mathbf{q}) = \sum_{i=1}^m \alpha(\mathbf{q}, \mathbf{k}_i)\mathbf{v}_i\\ &a(\mathbf{q}, \mathbf{k}_i) = \mathbf{w}_v^\top\tanh(\mathbf{q}\mathbf{W}_q + \mathbf{k}_i\mathbf{W}_k)\\ &\alpha(\mathbf{q}, \mathbf{k}_i) = \mathrm{softmax}\left(a(\mathbf{q}, \mathbf{k}_i)\right) \end{aligned} \]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 加性注意力
class AdditiveAttention(nn.Module):
    def __init__(self, key_size, query_size, atten_size):
        super().__init__()
        self.W_k = nn.Linear(key_size, atten_size, bias=False)
        self.W_q = nn.Linear(query_size, atten_size, bias=False)
        self.w_v = nn.Linear(atten_size, 1, bias=False)
        self.tanh = nn.Tanh()
        
    def forward(self, keys, values, queries):
        '''
        keys: [batch_size, src_len, key_size]
        values: [batch_size, src_len, value_size]
        queries: [batch_size, query_size]
        '''
        keys = self.W_k(keys)	# [batch_size, src_len, atten_size]
        queries = self.W_q(queries)	# [batch_size, atten_size]
        scores = self.w_v(self.tanh(keys + queries.unsqueeze(1))).squeeze(-1)	# [batch_size, src_len]
        atten_weights = torch.softmax(scores, dim=-1)	# [batch_size, src_len]
        context = torch.bmm(atten_weights.unsqueeze(1), values).squeeze(1)	# [batch_size, value_size]
        return context, atten_weights

解码器

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
# 解码器
class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, dec_hidden_size, enc_hidden_size, attention):
        super().__init__()
        self.attention = attention
        self.trg_vocab_size = vocab_size
        
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + enc_hidden_size*2, dec_hidden_size, batch_first=True)
        self.fc = nn.Linear(dec_hidden_size + enc_hidden_size*2 + embed_size, vocab_size)
        
    def forward(self, input, hidden, encoder_outputs):
        """
        input: [batch_size]，每个样本每次输入一个token
        hidden: [batch_size, dec_hidden_size]
        encoder_outputs: [batch_size, src_len, enc_hidden_size*2]
        """
        input = input.unsqueeze(1)	# [batch_size, 1]
        embedded = self.embedding(input)	# [batch_size, 1, embed_size]
        
        # 计算注意力
        context, atten_weights = self.attention(
            keys=encoder_outputs,
            values=encoder_outputs,
            queries=hidden
        )
        # context: [batch_size, encoder_hidden_size*2]
        
        context = context.unsqueeze(1)	# [batch_size, 1, enc_hidden_size*2]
        
        rnn_input = torch.cat((embedded, context), dim=2)	# [batch_size, 1, embed_size + enc_hidden_size*2]
        hidden = hidden.unsqueeze(0)	# [1, batch_size, dec_hidden_size]
        
        output, hidden = self.rnn(rnn_input, hidden)
        # output: [batch_size, 1, dec_hidden_size]
        # hidden: [1, batch_size, dec_hidden_size]
        
        output = output.squeeze(1)	# [batch_size, dec_hidden_size]
        context = context.squeeze(1)	# [batch_size, enc_hidden_size*2]
        embedded = embedded.squeeze(1)	# [batch_size, embed_size]
        
        prediction = self.fc(
            torch.cat((output, context, embedded), dim=1)
        )	# [batch_size, trg_vocab_size]
        
        return prediction, hidden.squeeze(0), atten_weights

Seq2Seq

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
class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
    def forward(self, src, src_len, trg, teacher_forcing_ratio=0.5):
        # 训练
        """
        src: [batch_size, src_len]
        trg: [batch_size, trg_len]
        """
        batch_size = src.size(0)
        trg_len = trg.size(1)
        
        encoder_outputs, hidden = self.encoder(src, src_len)
        input = trg[:, 0]	# 初始输入是`<SOS>`标记
        
        # 存储输出，提前构建张量
        outputs = torch.zeros(batch_size, trg_len, self.decoder.trg_vocab_size).to(device)
        
        for t in range(1, trg_len):
            output, hidden, atten_weights = self.decoder(input, hidden, encoder_outputs)
            outputs[:, t] = output	# outputs从[:, 1]开始，outputs[:, 0]并未被改变
            
            # 是否使用教师强制
            teacher_force = torch.rand(1) < teacher_forcing_ratio
            top1 = output.argmax(1)	# [batch_size]
            input = trg[:, t] if teacher_force else top1
        
        return outputs
    
    def predict(self, src, src_len, max_len=50):
        # 推理模式
        self.eval()
        with torch.no_grad():
            encoder_outputs, hidden = self.encoder(src, src_len)
            
            # 初始输入是`<SOS>`
            input = torch.zeros(src.size(0), 1).fill_(SOS_IDX).long().to(self.device)
            
            outputs = []
            atten_weights = []
            
            for t in range(1, max_len):
                output, hidden, attn = self.decoder(input.squeeze(1), hidden, encoder_outputs)
                outputs.append(output)
                atten_weights.append(attn)
                
                # 获取预测
                top1 = output.argmax(1)	# [batch_size]
                input = top1.unsqueeze(1)	# [batch_size, 1]
                
                # 遇到`<EOS>`停止
                if(input == EOS_IDX).all():
                    break
                
            outputs = torch.stack(outputs, dim=1)	# batch中的同一个样本的拼接在一起
            atten_weights = torch.stack(atten_weights, dim=1)
            return outputs, atten_weights

使用注意力机制的Seq2Seq进行机器翻译

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
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.utils.rnn import pad_sequence, pad_packed_sequence, pack_padded_sequence
from torch.utils.data import Dataset, DataLoader
import math
from collections import Counter
import re
from tqdm import tqdm

PAD_IDX = 0
SOS_IDX = 1
EOS_IDX = 2
UNK_IDX = 3

# 编码器
class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, enc_hidden_size, num_layers=1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, enc_hidden_size, num_layers=num_layers, bidirectional=True, batch_first=True)
        self.fc = nn.Linear(enc_hidden_size*2, enc_hidden_size)
    
    def forward(self, src, src_len):
        embedded = self.embedding(src)
        packed = pack_padded_sequence(embedded, src_len.cpu(), batch_first=True, enforce_sorted=True)
        outputs, hidden = self.rnn(packed)
        outputs, _ = pad_packed_sequence(outputs, batch_first=True)
        
        # 合并双向最后层的隐藏状态
        hidden = torch.tanh(self.fc(torch.cat((hidden[-1], hidden[-2]), dim=1)))
        # hidden: [Batch_size, enc_hidden_size]
        return outputs, hidden
    
class AdditiveAttention(nn.Module):
    def __init__(self, key_size, query_size, atten_size):
        """
        key_size: 编码器输出的特征维度
        query_size: 解码器隐状态的特征维度
        atten_size: 注意力空间维度
        """
        super().__init__()
        self.W_k = nn.Linear(key_size, atten_size, bias=False)
        self.W_q = nn.Linear(query_size, atten_size, bias=False)
        self.w_v = nn.Linear(atten_size, 1, bias=False)
        self.tanh = nn.Tanh()
        
    def forward(self, keys, values, queries):
        """
        keys(K): [Batch_size, seq_len, key_size] 编码器的输出序列（双向）
        values(V): [Batch_size, seq_len, value_size] 同Keys
        queries(Q): [Batch_size, queryp_size] 解码器当前隐状态
        """
        K = self.W_k(keys)
        Q = self.W_q(queries).unsqueeze(1)

        scores = self.w_v(self.tanh(Q + K)).squeeze(-1)
        alpha = torch.softmax(scores, dim=-1)

        context = torch.bmm(alpha.unsqueeze(1), values).squeeze(1)
        return context, alpha
    
# 解码器
class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, dec_hidden_size, enc_hidden_size, attention):
        super().__init__()
        self.attention = attention
        self.trg_vocab_size = vocab_size
        
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + enc_hidden_size*2, dec_hidden_size, batch_first=True)
        self.fc = nn.Linear(dec_hidden_size + enc_hidden_size*2 + embed_size, vocab_size)

    def forward(self, input, hidden, encoder_outputs):
        """
        input: [batch_size]
        hidden: [batch_size, dec_hidden_size]
        encoder_outputs: [batch_size, src_len, enc_hidden_size*2]
        """
        input = input.unsqueeze(1)  # [Batch_size, 1]
        embedded = self.embedding(input)    #[batch_size, 1, embed_size]

        # 计算注意力
        context, atten_weights = self.attention(
            keys=encoder_outputs,
            values=encoder_outputs,
            queries=hidden
        )
        # context: [batch_size, enc_hidden_size*2]
        context = context.unsqueeze(1)  # [batch_size, 1, enc_hidden_size*2]

        # RNN输入=嵌入向量+上下文
        rnn_input = torch.cat((embedded, context), dim=2)   # [batch_size, 1, embed_size+enc_hidden_size*2]
        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
        # output: [batch_size, 1, dec_hidden_size]
        # hidden: [1, batch_size, dec_hidden_size]

        # 最终预测
        output = output.squeeze(1)  # [batch_size, dec_hidden_size]
        context = context.squeeze(1)    # [batch_size, enc_hidden_size*2]
        embedded = embedded.squeeze(1)  # [batch_size, embed_size]

        prediction = self.fc(torch.cat(
            (output, context, embedded), dim=1
        ))
        return prediction, hidden.squeeze(0), atten_weights

# Seq2Seq
class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, src_len, trg, teacher_forcing_ratio=0.5):
        # src: [batch_size, src_len]
        # trg: [batch_size, trg_len]
        batch_size = src.size(0)
        trg_len = trg.size(1)

        # 存储输出
        outputs = torch.zeros(batch_size, trg_len, self.decoder.trg_vocab_size).to(self.device)

        # 编码
        encoder_outputs, hidden = self.encoder(src, src_len)

        # 初始输入是`<sos>`标记
        input = trg[:, 0]

        # 解码循环
        for i in range(1, trg_len):
            output, hidden, _ = self.decoder(input, hidden, encoder_outputs)
            outputs[:, i] = output

            # 决定是否使用教师强制
            teacher_force = torch.rand(1) < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[:, i] if teacher_force else top1

        return outputs
    
    def predict(self, src, src_len, max_len=50):
        # 推理模式
        self.eval()
        with torch.no_grad():
            # 编码
            encoder_outputs, hidden = self.encoder(src, src_len)
            # 初始输入是'<sos>'
            input = torch.ones(src.size(0), 1).fill_(SOS_IDX).long().to(self.device)

            outputs = []
            attn_weights = []

            for t in range(1, max_len):
                output, hidden, attn = self.decoder(input.squeeze(1), hidden, encoder_outputs)
                outputs.append(output)
                attn_weights.append(attn)

                # 获取预测词
                top1 = output.argmax(1)
                input = top1.unsqueeze(1)

                # 遇到'<eos>'停止
                if(input == EOS_IDX).all():
                    break
            
            outputs = torch.stack(outputs, 1)
            attn_weights = torch.stack(attn_weights, 1)
            return outputs, attn_weights


class TextDataset(Dataset):
    def __init__(self, en_list, cn_list):
        self.cn_data = cn_list
        self.en_data = en_list

    def __len__(self):
        return len(self.en_data)
    
    def __getitem__(self, index):
        return (
            torch.tensor(self.en_data[index], dtype=torch.long),
            torch.tensor(self.cn_data[index], dtype=torch.long),
            torch.tensor(len(self.en_data[index]), dtype=torch.long)
        )
    
    
    
def collate_fn(batch):
    batch.sort(key=lambda x: x[2], reverse=True)
    en_sequence, cn_sequence, en_lengths = zip(*batch)
    padded_en_sequence = pad_sequence(en_sequence, batch_first=True, padding_value=PAD_IDX)
    padded_cn_sequence = pad_sequence(cn_sequence, batch_first=True, padding_value=PAD_IDX)
    en_lengths = torch.tensor(en_lengths)

    return padded_en_sequence, padded_cn_sequence, en_lengths



def en_tokenize(text):
    text = text.lower()
    text_list = re.split('(\W)', text)
    return [text.strip() for text in text_list if text.strip()]

def cn_tokenize(text):
    return list(text)


def build_vocab(tokenized_text, vocab_limit=800000):
    if isinstance(tokenized_text[0], list):
        tokenized_data = [text for text_list in tokenized_text for text in text_list]
    else:
        tokenized_data = tokenized_text
    word_count = Counter(tokenized_data)
    vocab = sorted(word_count, key=word_count.get, reverse=True)[: vocab_limit]
    vocab = ['<PAD>', '<SOS>', '<EOS>', '<UNK>'] + vocab
    word_to_idx = {word: i for i, word in enumerate(vocab)}
    return word_to_idx, vocab



def train(model, iterator, device, optimizer, loss_function):
    model.train()
    total_loss = 0

    for src, trg, src_len in tqdm(iterator):
        src, trg = src.to(device), trg.to(device)
        optimizer.zero_grad()
        output = model(src, src_len, trg)   # [batch_size, trg_len, trg_vocab_size]

        loss = loss_function(output[:, 1:].reshape(-1, output.shape[-1]), trg[:, 1:].reshape(-1))
        total_loss += loss.item()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
    
    return total_loss / len(iterator)


with open('./中英翻译数据集/train.zh', 'r') as f:
    cn_data_1 = f.readlines()[:1000]
    cn_data = [text.strip() for text in cn_data_1]

with open('./中英翻译数据集/train.en', 'r') as f:
    en_data_1 = f.readlines()[:1000]
    en_data = [text.strip() for text in en_data_1]
    
    
cn_data_tokenized = [cn_tokenize(text) for text in cn_data]
en_data_tokenized = [en_tokenize(text) for text in en_data]

cn_word_to_idx, cn_vocab = build_vocab(cn_data_tokenized)
cn_idx_to_word = {i: word for i, word in enumerate(cn_vocab)}
cn_vocab_size = len(cn_vocab)

en_word_to_idx, en_vocab = build_vocab(en_data_tokenized)
en_idx_to_word = {i: word for i, word in enumerate(en_vocab)}
en_vocab_size = len(en_vocab)

# 英译中，中文开头需要为`<SOS>`，结尾需要为`<EOS>`
cn_data_tokenized = [['<SOS>'] + text_list + ['<EOS>'] for text_list in cn_data_tokenized]

cn_data_idx = [[cn_word_to_idx[word] if word in cn_word_to_idx else UNK_IDX for word in text_list] for text_list in cn_data_tokenized]
en_data_idx = [[en_word_to_idx[word] if word in en_word_to_idx else UNK_IDX for word in text_list] for text_list in en_data_tokenized]


dataset = TextDataset(en_data_idx, cn_data_idx)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)


enc_hidden_size = 512
dec_hidden_size = 512
enc_emb_size = 256
dec_emb_size = 256
atten_size = 256
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
attention = AdditiveAttention(
    key_size = enc_hidden_size*2,
    query_size=dec_hidden_size,
    atten_size=atten_size
)

encoder = Encoder(
    vocab_size=en_vocab_size,
    embed_size=enc_emb_size,
    enc_hidden_size=enc_hidden_size,
    num_layers=3
)

decoder = Decoder(
    vocab_size=cn_vocab_size,
    dec_hidden_size=dec_hidden_size,
    enc_hidden_size=enc_hidden_size,
    embed_size=dec_emb_size,
    attention=attention
)

model = Seq2Seq(encoder, decoder, device).to(device)


optimizer = optim.Adam(model.parameters())
loss_function = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
# 损失函数忽略<PAD>，防止填充部分被学习

num_epochs = 5
for epoch in range(num_epochs):
    Loss = train(model, dataloader, device, optimizer, loss_function)
    print(f'Epoch {epoch + 1}, Loss {Loss:.4f}')

注意：注意力机制中的K、Q等，会随着batch_size的增大而增大，随着序列长度的增大而增大，有可能会训练中突然爆显存，因此要控制合适的batch_size、序列长度

深度学习 > 理论基础
#人工智能 #神经网络 #深度学习 #NLP
使用注意力机制的Seq2Seq
https://blog.shinebook.net/2025/04/19/人工智能/理论基础/深度学习/使用注意力机制的Seq2Seq/
作者
X
发布于
2025年4月19日
许可协议