大模型基础补全计划(六)---带注意力机制的seq2seq实例与测试(Bahdanau Attention)

PS：要转载请注明出处，本人版权所有。

PS: 这个只是基于《我自己》的理解，

如果和你的原则及想法相冲突，请谅解，勿喷。

环境说明

无

前言

本文是这个系列第六篇，它们是：

《大模型基础补全计划(一)---重温一些深度学习相关的数学知识》 https://www.cnblogs.com/Iflyinsky/p/18717317
《大模型基础补全计划(二)---词嵌入(word embedding) 》 https://www.cnblogs.com/Iflyinsky/p/18775451
《大模型基础补全计划(三)---RNN实例与测试》 https://www.cnblogs.com/Iflyinsky/p/18967569
《大模型基础补全计划(四)---LSTM的实例与测试(RNN的改进)》 https://www.cnblogs.com/Iflyinsky/p/19091089
《大模型基础补全计划(五)---seq2seq实例与测试(编码器、解码器架构)》 https://www.cnblogs.com/Iflyinsky/p/19150535

本文，介绍一下注意力机制，并在上文的机翻模型seq2seq的实例中添加一个简单的注意力机制，并看看模型效果是否有提升。

注意力机制（Bahdanau Attention）

举一个例子：在日常生活中，比如我们看一幅黑白画（画中有一个红色的苹果，其他的都是黑白的物体，例如香蕉），这个时候我们无意识的看一眼画，很有可能第一个关注的就是这个红色的苹果，但是我有意识的控制眼睛集中去看香蕉，这个时候我关注的就是香蕉。

在上面的例子中，我们的注意力，最开始是无意识的看苹果，后面有意识的注意香蕉，这里面的区别就是我们在这个动作里面加了：意识。当加了意识后，我们就可以有选择的根据条件来关注这幅画的我想关注的地方。

然后我们可以对上面的现象进行建模：$R=Attention(Q,K)*V$，这里我们将Attention当作意识，V当作黑白画的特征，Q是画中是什么？K是V的标签（你可以把K当作是V有关联的部分，不同的K，对应的不同的V），如果没有Attention，R就是苹果，有了Attention，R就可以是香蕉。

我们回头想一想上一篇文的seq2seq中，我们的encoder的output是最后一层rnn的所有时间步的隐藏状态（num_steps,batch_size,num_hiddens），这里包含了我们的序列数据在不同时间步的特征变化，当我们在做decoder的时候，我们是拿着这个encoder的最后一层rnn的最后一个时间步的隐藏状态（1,batch_size,num_hiddens）来作为context的，是一个固定的值，这样有几个问题：

对于长序列来说，context可能丢失信息。
我们从固定context中解码信息，导致了我们对序列在特定解码步骤中，对context关注重点是一样的。

针对上面seq2seq的问题，Bahdanau设计了一种模型，可以解决我们遇到的问题，其定义如下：$$c_{t'} = \sum_{t=1}^{T} \alpha(s_{t'-1}, h_{t})h_{t}$$，看公式我们可以知道，这里定义了Q（decoder的上一次隐藏态$s_{t'-1}$）/K（encoder的output的部分$h_{t}$）/V（encoder的output的部分$h_{t}$）三个概念，含义就是通过Q+K来计算一个权重矩阵W（通过softmax归一化），然后然后将W和V进行计算，得到了我们通过W关注到的新的$V_{new}$，这里的W就是我们的注意力矩阵，代表我们关注V中的哪些部分。整个计算过程，就相当于我们生成了新成context具备了注意力机制。

带注意力机制的seq2seq 英文翻译中文的实例

下面的代码和上一篇文章的代码只有decoder部分有比较大的差别，其他的基本类似。如dataset部分的内容，请参考上一篇文章。

seq2seq完整代码如下

import os
import random
import torch
import math
from torch import nn
from torch.nn import functional as F
import numpy as np
import time
import visdom
import collectionsimport dataset
class Accumulator:"""在n个变量上累加"""def __init__(self, n):"""Defined in :numref:`sec_softmax_scratch`"""self.data = [0.0] * ndef add(self, *args):self.data = [a + float(b) for a, b in zip(self.data, args)]def reset(self):self.data = [0.0] * len(self.data)def __getitem__(self, idx):return self.data[idx]class Timer:"""记录多次运行时间"""def __init__(self):"""Defined in :numref:`subsec_linear_model`"""self.times = []self.start()def start(self):"""启动计时器"""self.tik = time.time()def stop(self):"""停止计时器并将时间记录在列表中"""self.times.append(time.time() - self.tik)return self.times[-1]def avg(self):"""返回平均时间"""return sum(self.times) / len(self.times)def sum(self):"""返回时间总和"""return sum(self.times)def cumsum(self):"""返回累计时间"""return np.array(self.times).cumsum().tolist()class Encoder(nn.Module):"""编码器-解码器架构的基本编码器接口"""def __init__(self, **kwargs):# 调用父类nn.Module的构造函数，确保正确初始化super(Encoder, self).__init__(**kwargs)def forward(self, X, *args):# 抛出未实现错误，意味着该方法需要在子类中具体实现raise NotImplementedErrorclass Decoder(nn.Module):"""编码器-解码器架构的基本解码器接口Defined in :numref:`sec_encoder-decoder`"""def __init__(self, **kwargs):# 调用父类nn.Module的构造函数，确保正确初始化super(Decoder, self).__init__(**kwargs)def init_state(self, enc_outputs, *args):# 抛出未实现错误，意味着该方法需要在子类中具体实现raise NotImplementedErrordef forward(self, X, state):# 抛出未实现错误，意味着该方法需要在子类中具体实现raise NotImplementedErrorclass EncoderDecoder(nn.Module):"""编码器-解码器架构的基类Defined in :numref:`sec_encoder-decoder`"""def __init__(self, encoder, decoder, **kwargs):# 调用父类nn.Module的构造函数，确保正确初始化super(EncoderDecoder, self).__init__(**kwargs)# 将传入的编码器实例赋值给类的属性self.encoder = encoder# 将传入的解码器实例赋值给类的属性self.decoder = decoderdef forward(self, enc_X, dec_X, enc_X_valid_len, *args):# 调用编码器的前向传播方法，处理输入的编码器输入数据enc_Xenc_outputs = self.encoder(enc_X, *args)# 调用解码器的init_state方法，根据编码器的输出初始化解码器的状态dec_state = self.decoder.init_state(enc_outputs, enc_X_valid_len)# 调用解码器的前向传播方法，处理输入的解码器输入数据dec_X和初始化后的状态return self.decoder(dec_X, dec_state)def masked_softmax(X, valid_lens):  #@save"""执行带掩码的 Softmax 操作。参数:X (torch.Tensor): 待 Softmax 的张量，通常是注意力机制中的“分数”（scores）。其形状通常为 (批量大小, 查询数量/序列长度, 键值对数量/序列长度)。valid_lens (torch.Tensor): 序列的有效长度张量。形状可以是 (批量大小,) 或 (批量大小, 键值对数量)。用于指示每个序列的哪个部分是有效的（非填充）。返回:torch.Tensor: 经过 Softmax 归一化且填充部分被忽略的概率分布张量。"""# 辅助函数：创建一个序列掩码，并用特定值覆盖被掩码（填充）的元素def _sequence_mask(X, valid_len, value=0):"""根据有效长度（valid_len）创建掩码，并应用于张量 X。参数:X (torch.Tensor): 形状为 (批量大小 * 查询数量, 最大长度) 的张量。valid_len (torch.Tensor): 形状为 (批量大小 * 查询数量,) 的有效长度向量。value (float): 用于替换被掩码元素的填充值。返回:torch.Tensor: 被填充值覆盖后的张量 X。"""# 获取序列的最大长度（张量的第二个维度）maxlen = X.size(1)# 核心掩码逻辑：# 1. torch.arange((maxlen), ...) 创建一个从 0 到 maxlen-1 的序列（代表时间步索引）# 2. [None, :] 使其形状变为 (1, maxlen)，用于广播# 3. valid_len[:, None] 使有效长度形状变为 (批量大小 * 查询数量, 1)，用于广播# 4. < 比较操作：当索引 < 有效长度时，结果为 True（有效元素），否则为 False（填充元素）mask = torch.arange((maxlen), dtype=torch.float32,device=X.device)[None, :] < valid_len[:, None]# 逻辑非 ~mask 得到填充部分的掩码（True 表示填充部分）# 使用填充值（value，通常是 -1e6）覆盖填充元素X[~mask] = valuereturn X# 1. 处理无需掩码的情况if valid_lens is None:# 如果未提供有效长度，则执行标准 Softmaxreturn nn.functional.softmax(X, dim=-1)# 2. 处理需要掩码的情况else:# 备份原始形状，用于后续重塑shape = X.shape# 统一 valid_lens 的形状，使其与 X 的前两个维度相匹配if valid_lens.dim() == 1:# 适用于批量中每个序列只有一个有效长度的情况（例如，K-V 序列是等长的）# 将 valid_lens 重复 shape[1] 次，匹配 X 的查询/序列长度维度valid_lens = torch.repeat_interleave(valid_lens, shape[1])else:# 适用于每个查询-键值对的有效长度都不同的情况# 将 2D 张量展平为 1D 向量valid_lens = valid_lens.reshape(-1)# 预处理 Softmax 输入：将 X 调整为 2D 矩阵 (批量*查询数量, 键值对数量)# 并在最后一个轴（Softmax 轴）上，用一个非常大的负值替换被掩码的元素# Softmax 时 exp(-1e6) 趋近于 0，从而忽略填充部分。X = _sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)# 对经过掩码处理的 X 执行 Softmax# 结果张量 X 被重塑回原始的 3D 形状 (批量大小, 查询数量, 键值对数量)# 并在最后一个维度（dim=-1）上进行归一化，得到注意力权重return nn.functional.softmax(X.reshape(shape), dim=-1)class AdditiveAttention(nn.Module):  #@save"""加性注意力（Additive Attention）模块。通过将 Query 和 Key 投影到相同的维度后相加，再通过 tanh 激活和线性层计算注意力分数。公式核心：score(Q, K) = w_v^T * tanh(W_q*Q + W_k*K)"""def __init__(self, num_hiddens, dropout, **kwargs):"""初始化加性注意力模块。参数:num_hiddens (int): 隐藏层维度，Q 和 K 投影后的维度。dropout (float): Dropout 率。"""super(AdditiveAttention, self).__init__(**kwargs)# W_k：将 Key (K) 向量投影到 num_hiddens 维度的线性层# 使用 nn.LazyLinear 延迟初始化，直到第一次 forward 传入 K 的维度self.W_k = nn.LazyLinear(num_hiddens, bias=False)# W_q：将 Query (Q) 向量投影到 num_hiddens 维度的线性层# 使用 nn.LazyLinear 延迟初始化self.W_q = nn.LazyLinear(num_hiddens, bias=False)# w_v：将激活后的特征向量 (W_q*Q + W_k*K) 投影成一个标量分数（维度为 1）# 使用 nn.LazyLinear 延迟初始化self.w_v = nn.LazyLinear(1, bias=False)# Dropout 层，用于防止过拟合self.dropout = nn.Dropout(dropout)def forward(self, queries, keys, values, valid_lens):"""执行前向传播计算。参数:queries (torch.Tensor): 查询向量 Q。形状通常为 (批量大小, 查询数量, 查询维度)。keys (torch.Tensor): 键向量 K。形状通常为 (批量大小, 键值对数量, 键维度)。values (torch.Tensor): 值向量 V。形状通常为 (批量大小, 键值对数量, 值维度)。valid_lens (torch.Tensor): 键值对序列的有效长度，用于掩盖填充部分。返回:torch.Tensor: 注意力加权后的值向量。形状为 (批量大小, 查询数量, 值维度)。"""# 1. 线性变换：分别对 Q 和 K 进行投影queries, keys = self.W_q(queries), self.W_k(keys)# 2. 维度扩展与相加（Attention Scoring 的核心步骤）# queries.unsqueeze(2): 形状从 (批量大小, 查询数量, num_hiddens) #                       变为 (批量大小, 查询数量, 1, num_hiddens)。# keys.unsqueeze(1): 形状从 (批量大小, 键值对数量, num_hiddens) #                     变为 (批量大小, 1, 键值对数量, num_hiddens)。# 两个张量通过广播机制相加，得到 features，形状为：# (批量大小, 查询数量, 键值对数量, num_hiddens)features = queries.unsqueeze(2) + keys.unsqueeze(1)# 3. 激活函数：应用 tanh 激活（加性注意力机制的要求）features = torch.tanh(features)# 4. 投影到标量分数# self.w_v(features): 将 features 的最后一个维度（num_hiddens）投影成 1。# scores.squeeze(-1): 移除最后一个单维度 (1)，得到最终的注意力分数张量。# 形状为：(批量大小, 查询数量, 键值对数量)scores = self.w_v(features).squeeze(-1)# 5. 归一化（Softmax）：使用带掩码的 Softmax 得到注意力权重# 填充部分的得分会被设置为一个极小的负值，Softmax 后权重趋近于 0。self.attention_weights = masked_softmax(scores, valid_lens)# 6. 加权求和# torch.bmm: 批量矩阵乘法 (Batch Matrix Multiplication)。# 将 [注意力权重] (批量, Q数量, K数量) 与 [值向量] (批量, K数量, V维度) 相乘# 得到最终的注意力输出，形状为：(批量大小, 查询数量, 值维度)# 在 BMM 之前，对注意力权重应用 Dropout。return torch.bmm(self.dropout(self.attention_weights), values)#@save
class Seq2SeqEncoder(Encoder):"""用于序列到序列学习的循环神经网络编码器"""def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,dropout=0, **kwargs):super(Seq2SeqEncoder, self).__init__(**kwargs)# 嵌入层self.embedding = nn.Embedding(vocab_size, embed_size)self.rnn = nn.GRU(embed_size, num_hiddens, num_layers,dropout=dropout)# self.lstm = nn.LSTM(embed_size, num_hiddens, num_layers)def forward(self, X, *args):# 输入X.shape = (batch_size,num_steps)# 输出'X'的形状：(batch_size,num_steps,embed_size)X = self.embedding(X)# 在循环神经网络模型中，第一个轴对应于时间步X = X.permute(1, 0, 2)# 如果未提及状态，则默认为0output, state = self.rnn(X)# output : 这个返回值是所有时间步的隐藏状态序列# output的形状:(num_steps,batch_size,num_hiddens)# hn (hidden) : 这是每一层rnn的最后一个时间步的隐藏状态# state的形状:(num_layers,batch_size,num_hiddens)return output, stateclass AttentionDecoder(Decoder):  #@save"""The base attention-based decoder interface."""def __init__(self):super().__init__()@propertydef attention_weights(self):raise NotImplementedErrorclass Seq2SeqAttentionDecoder(AttentionDecoder):def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,dropout=0):super().__init__()self.attention = AdditiveAttention(num_hiddens, dropout)self.embedding = nn.Embedding(vocab_size, embed_size)self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers,dropout=dropout)self.dense = nn.LazyLinear(vocab_size)# self.apply(d2l.init_seq2seq)def init_state(self, enc_outputs, enc_valid_lens):# Shape of outputs: (num_steps, batch_size, num_hiddens).# Shape of hidden_state: (num_layers, batch_size, num_hiddens)outputs, hidden_state = enc_outputsreturn (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)def forward(self, X, state):# Shape of enc_outputs: (batch_size, num_steps, num_hiddens).# Shape of hidden_state: (num_layers, batch_size, num_hiddens)enc_outputs, hidden_state, enc_valid_lens = state# Shape of the output X: (num_steps, batch_size, embed_size)X = self.embedding(X).permute(1, 0, 2)outputs, self._attention_weights = [], []for x in X:# Shape of query: (batch_size, 1, num_hiddens)query = torch.unsqueeze(hidden_state[-1], dim=1)# Shape of context: (batch_size, 1, num_hiddens)context  = self.attention(query, enc_outputs, enc_outputs, enc_valid_lens)# Concatenate on the feature dimensionx = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)# Reshape x as (1, batch_size, embed_size + num_hiddens)out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)outputs.append(out)self._attention_weights.append(self.attention.attention_weights)# After fully connected layer transformation, shape of outputs:# (num_steps, batch_size, vocab_size)outputs = self.dense(torch.cat(outputs, dim=0))return outputs.permute(1, 0, 2), [enc_outputs, hidden_state,enc_valid_lens]@propertydef attention_weights(self):return self._attention_weightsdef try_gpu(i=0):"""如果存在，则返回gpu(i)，否则返回cpu()Defined in :numref:`sec_use_gpu`"""if torch.cuda.device_count() >= i + 1:return torch.device(f'cuda:{i}')return torch.device('cpu')def sequence_mask(X, valid_len, value=0):"""在序列中屏蔽不相关的项"""maxlen = X.size(1)mask = torch.arange((maxlen), dtype=torch.float32,device=X.device)[None, :] < valid_len[:, None]X[~mask] = valuereturn Xclass MaskedSoftmaxCELoss(nn.CrossEntropyLoss):"""带遮蔽的softmax交叉熵损失函数"""# pred的形状：(batch_size,num_steps,vocab_size)# label的形状：(batch_size,num_steps)# valid_len的形状：(batch_size,)def forward(self, pred, label, valid_len):weights = torch.ones_like(label)weights = sequence_mask(weights, valid_len)self.reduction='none'unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(pred.permute(0, 2, 1), label)weighted_loss = (unweighted_loss * weights).mean(dim=1)return weighted_lossdef grad_clipping(net, theta):  #@save"""裁剪梯度"""if isinstance(net, nn.Module):params = [p for p in net.parameters() if p.requires_grad]else:params = net.paramsnorm = torch.sqrt(sum(torch.sum((p.grad ** 2)) for p in params))if norm > theta:for param in params:param.grad[:] *= theta / normdef train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):"""训练序列到序列模型"""def xavier_init_weights(m):if type(m) == nn.Linear:nn.init.xavier_uniform_(m.weight)if type(m) == nn.GRU:for param in m._flat_weights_names:if "weight" in param:nn.init.xavier_uniform_(m._parameters[param])net.apply(xavier_init_weights)net.to(device)optimizer = torch.optim.Adam(net.parameters(), lr=lr)loss = MaskedSoftmaxCELoss()net.train()vis = visdom.Visdom(env=u'test1', server="http://127.0.0.1", port=8097)animator = visfor epoch in range(num_epochs):timer = Timer()metric = Accumulator(2)  # 训练损失总和，词元数量for batch in data_iter:#清零（reset）优化器中的梯度缓存optimizer.zero_grad()# x.shape = [batch_size, num_steps]X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]# bos.shape = batch_size 个 bos-idbos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],device=device).reshape(-1, 1)# dec_input.shape = (batch_size, num_steps)# 解码器的输入通常由序列的起始标志 bos 和目标序列（去掉末尾的部分 Y[:, :-1]）组成。dec_input = torch.cat([bos, Y[:, :-1]], 1)  # 强制教学# Y_hat的形状:(batch_size,num_steps,vocab_size)Y_hat, _ = net(X, dec_input, X_valid_len)l = loss(Y_hat, Y, Y_valid_len)l.sum().backward()      # 损失函数的标量进行“反向传播”grad_clipping(net, 1)num_tokens = Y_valid_len.sum()optimizer.step()with torch.no_grad():metric.add(l.sum(), num_tokens)if (epoch + 1) % 10 == 0:# print(predict('你是？'))# print(epoch)# animator.add(epoch + 1, )if epoch == 9:# 清空图表：使用空数组来替换现有内容vis.line(X=np.array([0]), Y=np.array([0]), win='train_ch8', update='replace')# _loss_val = l# _loss_val = _loss_val.cpu().sum().detach().numpy()vis.line(X=np.array([epoch + 1]),Y=[ metric[0] / metric[1]],win='train_ch8',update='append',opts={'title': 'train_ch8','xlabel': 'epoch','ylabel': 'loss','linecolor': np.array([[0, 0, 255]]),  # 蓝色线条})print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} 'f'tokens/sec on {str(device)}')torch.save(net.cpu().state_dict(), 'model_h.pt')  # [[6]]torch.save(net.cpu(), 'model.pt')  # [[6]]def predict_seq2seq(net, src_sentence, src_vocab, tgt_vocab, num_steps,device, save_attention_weights=False):"""序列到序列模型的预测"""# 在预测时将net设置为评估模式net.eval()src_tokens = src_vocab[src_sentence.lower().split(' ')] + [src_vocab['<eos>']]enc_valid_len = torch.tensor([len(src_tokens)], device=device)src_tokens = dataset.truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])# 添加批量轴enc_X = torch.unsqueeze(torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)enc_outputs = net.encoder(enc_X, enc_valid_len)dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)# 添加批量轴dec_X = torch.unsqueeze(torch.tensor([tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)output_seq, attention_weight_seq = [], []for _ in range(num_steps):Y, dec_state = net.decoder(dec_X, dec_state)# 我们使用具有预测最高可能性的词元，作为解码器在下一时间步的输入dec_X = Y.argmax(dim=2)pred = dec_X.squeeze(dim=0).type(torch.int32).item()# 保存注意力权重（稍后讨论）if save_attention_weights:attention_weight_seq.append(net.decoder.attention_weights[0].reshape(num_steps).cpu())# 一旦序列结束词元被预测，输出序列的生成就完成了if pred == tgt_vocab['<eos>']:breakoutput_seq.append(pred)return ' '.join(tgt_vocab.to_tokens(output_seq)), attention_weight_seqdef bleu(pred_seq, label_seq, k):  #@save"""计算BLEU"""pred_tokens, label_tokens = pred_seq.split(' '), [i for i in label_seq]len_pred, len_label = len(pred_tokens), len(label_tokens)score = math.exp(min(0, 1 - len_label / len_pred))for n in range(1, k + 1):num_matches, label_subs = 0, collections.defaultdict(int)for i in range(len_label - n + 1):label_subs[' '.join(label_tokens[i: i + n])] += 1for i in range(len_pred - n + 1):if label_subs[' '.join(pred_tokens[i: i + n])] > 0:num_matches += 1label_subs[' '.join(pred_tokens[i: i + n])] -= 1score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))return scorefrom matplotlib import pyplot as plt
import matplotlib
# from matplotlib_inline import backend_inline
def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),cmap='Reds'):"""显示矩阵的热图（Heatmaps）。这个函数旨在以子图网格的形式绘制多个矩阵，通常用于可视化注意力权重等。参数:matrices (numpy.ndarray 或 torch.Tensor 数组): 一个四维数组，形状应为 (num_rows, num_cols, height, width)。其中，num_rows 和 num_cols 决定了子图网格的布局，height 和 width 是每个热图（即每个矩阵）的维度。xlabel (str): 所有最底行子图的 x 轴标签。ylabel (str): 所有最左列子图的 y 轴标签。titles (list of str, optional): 一个包含 num_cols 个标题的列表，用于设置每一列子图的标题。默认 None。figsize (tuple, optional): 整个图形（figure）的大小。默认 (2.5, 2.5)。cmap (str, optional): 用于绘制热图的颜色映射（colormap）。默认 'Reds'。"""# 导入所需的 matplotlib 模块，确保图形在 Jupyter/IPython 环境中正确显示为 SVG 格式# （假设在包含这个函数的环境中已经导入了 matplotlib 的 backend_inline）# backend_inline.set_matplotlib_formats('svg')matplotlib.use('TkAgg')# 从输入的 matrices 形状中解构出子图网格的行数和列数# 假设 matrices 的形状是 (num_rows, num_cols, height, width)num_rows, num_cols, _, _ = matrices.shape# 创建一个包含多个子图（axes）的图形（fig）# fig: 整个图形对象# axes: 一个 num_rows x num_cols 的子图对象数组fig, axes = plt.subplots(num_rows, num_cols, figsize=figsize,sharex=True,    # 所有子图共享 x 轴刻度sharey=True,    # 所有子图共享 y 轴刻度squeeze=False   # 即使只有一行或一列，也强制返回二维数组的 axes，方便后续循环)# 遍历子图的行和对应的矩阵行# i 是行索引, row_axes 是当前行的子图数组, row_matrices 是当前行的矩阵数组for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):# 遍历当前行中的子图和对应的矩阵# j 是列索引, ax 是当前的子图对象, matrix 是当前的待绘矩阵for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):# 使用 ax.imshow() 绘制热图# matrix.detach().numpy()：将 PyTorch Tensor 转换为 numpy 数组，并从计算图中分离（如果它是 Tensor）# cmap：指定颜色映射pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap)# --- 设置轴标签和标题 ---# 只有最底行 (i == num_rows - 1) 的子图才显示 x 轴标签if i == num_rows - 1:ax.set_xlabel(xlabel)# 只有最左列 (j == 0) 的子图才显示 y 轴标签if j == 0:ax.set_ylabel(ylabel)# 如果提供了标题列表，则设置当前列的子图标题（所有行共享列标题）if titles:ax.set_title(titles[j])# --- 添加颜色条（Colorbar） ---# 为整个图形添加一个颜色条，用于表示数值和颜色的对应关系# pcm: 之前绘制的第一个热图返回的 Colormap # ax=axes: 颜色条将参照整个子图网格进行定位和缩放# shrink=0.6: 缩小颜色条的高度/长度，使其只占图形高度的 60%fig.colorbar(pcm, ax=axes, shrink=0.6)plt.show()if __name__ == '__main__':embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1batch_size, num_steps = 64, 10lr, num_epochs, device = 0.005, 2000, try_gpu()# train_iter 每个迭代输出：(batch_size, num_steps)train_iter, src_vocab, tgt_vocab, source, target = dataset.load_data(batch_size, num_steps)encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hiddens, num_layers,dropout)decoder = Seq2SeqAttentionDecoder(len(tgt_vocab), embed_size, num_hiddens, num_layers,dropout)net = EncoderDecoder(encoder, decoder)is_train = Falseis_show = Falseif is_train:train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)elif is_show:state_dict = torch.load('model_h.pt')net.load_state_dict(state_dict)net.to(device)src_text = "Call us."translation, attention_weight_seq = predict_seq2seq(net, src_text, src_vocab, tgt_vocab, num_steps, device, True)# attention_weights = torch.eye(10).reshape((1, 1, 10, 10))# (num_rows, num_cols, height, width)print(f'translation={translation}')print(attention_weight_seq)stacked_tensor = torch.stack(attention_weight_seq, dim=0)stacked_tensor = stacked_tensor.unsqueeze(0).unsqueeze(0)show_heatmaps(stacked_tensor,xlabel='Attention weight', ylabel='Decode Step')else:state_dict = torch.load('model_h.pt')net.load_state_dict(state_dict)net.to(device)C = 0C1 = 0for i in range(2000):# print(source[i])# print(target[i])translation, attention_weight_seq = predict_seq2seq(net, source[i], src_vocab, tgt_vocab, num_steps, device)score = bleu(translation, target[i], k=2)if score > 0.0:C = C + 1if score > 0.8:C1 = C1 + 1print(f'{source[i]} => {translation}, bleu {score:.3f}')print(f'Counter(bleu > 0) = {C}')print(f'Valid-Counter(bleu > 0.8) = {C1}')

下面是encoder过程的简单分析：

将x通过nn.Embedding得到了(batch_size,num_steps,embed_size)的输入嵌入向量。
将嵌入向量传给nn.GRU，得到了两个输出，并返回：
- output，最后一层rnn的所有时间步的隐藏状态（num_steps,batch_size,num_hiddens）。
- h_n,所有rnn层的，最后一个时间步的隐藏状态（num_layers,batch_size,num_hiddens）。

下面是decoder过程的简单分析：

将decoder_x通过nn.Embedding得到了(batch_size,num_steps,embed_size)的输入嵌入向量。
将嵌入向量沿着num_steps进行单步运行，每一步经过Attention过程，得到最终的output，以及最后一个时间步的所有rnn层的h_n，每一步执行如下步骤：
- 将rnn最后一层的隐藏态作为Q（第一次Q是来自于encoder，后续都是decoder的每一次运行过程产生的隐藏态）
- 将encoder的output作为K,V，得到当前动态的上下文 context
- 将decoder_x_step 和 context进行组合，得到decoder_x_step_new
- 将decoder_x_step_new送入nn.GRU，得到当前时间步的output, h_t
- 将每一步的output收集起来作为输出，将h_t作为下一个时间步的Q循环起来
将所有的output经过nn.LazyLinear 映射为(num_steps, batch_size, vocab_size)，并和h_t返回

和原版本的seq2seq进行对比可知：

在原版中，我们的decoder依赖于一个固定的enc_outputs进行循环解码
在新版中，我们的decoder每次界面，都会有一个Q（第一次是来至于encoder，后续都是decoder的每一次运行过程产生的隐藏态）来计算enc_outputs的权重分数，然后根据权重分数得到一个动态的enc_outputs，这样可以让解码器每一步都关注enc_outputs中的不同的重点。

attention weight 的解释：

把Encoder Output（num_steps,1,num_hiddens）作为K,V
将Decoder的隐藏态h_t（1,1,num_hiddens）(初始值来自于Encoder的隐藏态h_t)作为Q，计算出当前step的attention_weight，其是一个softmax概率数据。
然后将attention_weight 与 V进行计算，代表模型当前关注EncoderOutput的那部分数据，得到新的Context

下面是训练和测试的一些结果

从上面的图可以看到，这个模型有一定的翻译效果，并且，比上一篇文章的模型效果要好一点。

此外，下面是我们翻译："Call us."-> "联系我们。" 的attention weight的可视化（带mask=3，在不同的decode step中权重变化。）

从上面的图可以知道，每一个decode step的注意力权重矩阵值都不一样，意味着，每一步解码的时候，关注的内容也不一样。

后记

本文引入了注意力机制，及注意力机制在seq2seq中，在应用注意力机制后，和原版的seq2seq的结论相比，模型效果有提升。

参考文献

https://zh.d2l.ai/chapter_recurrent-modern/encoder-decoder.html
https://zh.d2l.ai/chapter_recurrent-modern/seq2seq.html
https://d2l.ai/chapter_attention-mechanisms-and-transformers/bahdanau-attention.html

打赏、订阅、收藏、丢香蕉、硬币，请关注公众号（攻城狮的搬砖之路）