07. 注意力机制
一、什么是注意力机制
传统的 Seq2Seq 模型中，编码器在处理源句时，无论其长度如何，最终都只能将整句信息压缩为一个固定长度的上下文向量，用作解码器的唯一参考。这种设计存在两个显著问题：
- 信息压缩困难：固定向量难以完整表达长句或复杂语义，容易丢失关键信息。
- 缺乏动态感知：解码器在每一步生成中都只能依赖同一个上下文向量，难以根据不同位置的生成需要灵活提取信息。
为了解决上述问题，研究者引入了 Attention 机制。其核心思想是：解码器在生成目标序列的每一步时，不再依赖于一个静态的上下文向量，而是根据当前的解码状态，动态地从编码器各时间步的隐藏状态中选取最相关的信息，以辅助当前步的生成。这种机制赋予模型 “对齐” 能力，使其能够自动判断源句中哪些位置对当前的目标词更为重要，从而有效缓解信息瓶颈问题，提升生成质量与表达能力。
我们在终端中输入 pip 指令安装 PyTorch。默认是从国外的主站上下载，因此，我们可能会遇到网络不好的情况导致下载失败。我们可以在 pip 指令后通过 -f 指定国内镜像源下载。
pip install torch torchvision -f https://mirrors.aliyun.com/pytorch-wheels/cu132
国内常用的 pip 下载源列表：
- 阿里云 https://mirrors.aliyun.com/pypi/simple
- 清华大学 https://pypi.tuna.tsinghua.edu.cn/simple
- 中国科学技术大学 https://pypi.mirrors.ustc.edu.cn/simple
安装完 PyTorch 之后，我们可以通过如下代码查看 CUDA 是否可用：
import torch
print(f"PyTorch的版本：{torch.__version__}")
print(f"CUDA是否可用：{torch.cuda.is_available()}")
如果输出 CUDA 不可用，有可能你使用国内镜像只同步 PyPI 主站，不包含 PyTorch 官方 CUDA 专用包，
pip只会找到 CPU 版，找不到 CUDA 版。因此，我们还需要使用--extra-index-url https://download.pytorch.org/whl/cu132补充官方 CUDA 源 。
二、注意力机制的工作原理
注意力机制的核心思想，是 解码器在生成目标序列的每一步时，动态地从编码器的各个时间步的隐藏状态中提取当前所需的信息，而不再只依赖一个固定的上下文向量。这一机制通常通过以下 4 个关键步骤实现：
2.1、相关性计算
在目标序列生成的每一步，解码器都会计算当前时间步的隐藏状态与编码器各个时间步输出之间的相关性。这些相关性衡量了源句中每个位置对当前生成内容的重要程度，从而决定模型应将多少注意力分配给不同的源位置。相关性的计算依赖于特定的函数，通常被称为 注意力评分函数（attention scoring function）。
2.2、注意力权重计算
得到所有源位置的注意力评分后，使用 Softmax() 函数将其归一化为概率分布，作为注意力权重。得分越高的位置，其对应的权重越大，代表模型在当前生成中更关注该位置的信息。
2.3、上下文向量计算
将所有编码器输出按照注意力权重进行加权求和，得到一个上下文向量。这个向量就表示当前时间步，模型从源句中提取出的关键信息。
2.4、解码信息融合
在得到上下文向量后，解码器将其与当前时间步的隐藏状态进行拼接，以融合两者信息，最终通过线性变换和 Softmax，生成当前时间步目标词的概率分布。
三、注意力评分函数
注意力评分函数有多种实现方式。这里介绍三种常见的计算方法：点积评分（Dot）、通用点积评分（General）和 拼接评分（Concat）。它们虽然在结构上各有差异，但本质上都是用于衡量解码器当前隐藏状态与编码器各时间步隐藏状态之间的相关性，并据此分配注意力权重。
3.1、点积评分（Dot）
点积评分 是注意力机制中最简单、最直接的一种相关性评分方法。它通过 计算解码器当前时间步的隐藏状态与编码器每个时间步的隐藏状态的点积，来衡量二者之间的相关性。
\[score = h^{dec}_{t} \times h^{enc}
\]
其含义可以理解为：如果两个向量方向越一致（即越接近），它们的点积就越大，表示相关性越强，模型应当给予更多注意力。
3.2、通用点积评分（General）
通用点积评分 在点积的基础上引入了一个 可学习的权重矩阵 W，用于 先对编码器隐藏状态进行线性变换，再与解码器隐藏状态进行点积：
\[score = h^{dec}_{t} \times W \times h^{enc}_{t}
\]
该方法的设计动机主要是为了解决编码器和解码器隐藏状态维度不一致的问题。通过引入权重矩阵 W，不仅实现了维度对齐，也增强了模型对编码器输出的适应能力，从而提升了注意力机制的表达能力。
3.3、拼接评分（General）
拼接评分是一种表达能力更强的相关性评分方法。它的核心思想是：将解码器当前隐藏状态与编码器每个时间步的隐藏状态拼接为一个长向量，经过线性变换和非线性激活，最后用一个向量进行投影，得到最终打分值：
\[score = W_{2} \times tanh(W_{1} \times \begin{bmatrix} \mathbf{h^{enc}_{t}} \\ \mathbf{h^{dec}_{t}} \end{bmatrix})
\]
相比前两种方法，拼接评分方式在建模能力上更强。它不仅考虑了两个状态的数值关系，还引入非线性变换，能够捕捉更复杂的交互模式，更适合处理对齐关系复杂的任务场景。
四、中英翻译案例
我们可以从 阿里天池 中下载 中英翻译的数据集： https://tianchi.aliyun.com/dataset/174937 。
我们可以在终端中使用 pip 安装 pandas、sklearn 机器学习库和 ntlk 自然语言处理工具包。默认是从国外的主站上下载，因此，我们可能会遇到网络不好的情况导致下载失败。我们可以在 pip 指令后通过 -i 指定国内镜像源下载。
pip install scikit-learn jieba nltk -i https://pypi.tuna.tsinghua.edu.cn/simple
数据预处理：
import os
DATA_DRIECTORY = os.path.join(os.getcwd(), "assets", "data")
MODEL_DRIECTORY = os.path.join(os.getcwd(), "assets", "models")
UNK_TOKEN = "<unk>"
PAD_TOKEN = "<pad>"
SOS_TOKEN = "<sos>"
EOS_TOKEN = "<eos>"
SPECIAL_TOKENS = [UNK_TOKEN, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN]
import os
import pandas as pd
from nltk.tokenize import TreebankWordTokenizer
from sklearn.model_selection import train_test_split
def preprocess_data(file_path, test_size):
# 读取 TSV 格式的文件，
# sep 指定字段的分隔符，head=None 表示没有表头，
# secols 为要使用的列，可以使用索引，names 表示为列取的标签名
df = pd.read_csv(file_path, sep="\t", header=None, usecols=[0, 1], names=["en", "zh_cn"], encoding="utf-8").dropna()
train_df, test_df = train_test_split(df, test_size=test_size) # 划分训练集与测试集
# 构建中文词表，这里采用按字符切分
zh_cn_vacab_set = set([vacab for sentence in train_df["zh_cn"] for vacab in sentence if vacab.strip()])
zh_cn_vacab_list = SPECIAL_TOKENS + list(zh_cn_vacab_set) # 添加特殊标记
print("构建中文词表成功")
# 构建英文词表，基于正则匹配规则，将输入的英文句子拆分成包含单词和独立标点符号的词元列表（Token List）
tokenizer = TreebankWordTokenizer()
en_vacab_set = set([vacab for sentence in train_df["en"] for vacab in tokenizer.tokenize(sentence) if vacab.strip()])
en_vacab_list = SPECIAL_TOKENS + list(en_vacab_set) # 添加特殊标记
print("构建英文词表成功")
if not os.path.exists(DATA_DRIECTORY):
os.makedirs(DATA_DRIECTORY)
# 保存词表到文件中
file_path = os.path.join(DATA_DRIECTORY, "zh_cn_vocab.txt")
with open(file_path, "w", encoding="utf-8") as f:
f.write("\n".join(zh_cn_vacab_list))
print("保存中文词表到文件中成功")
file_path = os.path.join(DATA_DRIECTORY, "en_vocab.txt")
with open(file_path, "w", encoding="utf-8") as f:
f.write("\n".join(en_vacab_list))
print("保存英文词表到文件中成功")
# 构建词表到索引的映射
zh_cn_word2index = {word: index for index, word in enumerate(zh_cn_vacab_list)}
print("构建中文词表到索引的映射成功")
en_word2index = {word: index for index, word in enumerate(en_vacab_list)}
print("构建英文词表到索引的映射成功")
# 构建训练集
train_df["zh_cn"] = train_df["zh_cn"].apply(lambda sentence: get_word2index(zh_cn_word2index, list(sentence)))
train_df["en"] = train_df["en"].apply(lambda sentence: get_word2index(en_word2index, tokenizer.tokenize(sentence), True))
file_path = os.path.join(DATA_DRIECTORY, "translation_train.jsonl")
train_df.to_json(file_path, orient="records", lines=True) # 保存训练集
print("保存训练集数据成功")
# 构建测试集
test_df["zh_cn"] = test_df["zh_cn"].apply(lambda sentence: get_word2index(zh_cn_word2index, list(sentence)))
test_df["en"] = test_df["en"].apply(lambda sentence: get_word2index(en_word2index, tokenizer.tokenize(sentence), True))
# 保存测试集
file_path = os.path.join(DATA_DRIECTORY, "translation_test.jsonl")
test_df.to_json(file_path, orient="records", lines=True)
print("保存测试集数据成功")
def get_word2index(word2index, tokens, add_sos_eos=False):
if add_sos_eos:
tokens = [SOS_TOKEN] + tokens + [EOS_TOKEN]
return [word2index.get(token, SPECIAL_TOKENS.index(UNK_TOKEN)) for token in tokens]
if __name__ == "__main__":
test_size = 0.2
file_path = os.path.join(DATA_DRIECTORY, "cmn.txt")
preprocess_data(file_path, test_size)
创建数据集：
import torch
import pandas as pd
from torch.utils.data import Dataset
class TranslationDataset(Dataset):
def __init__(self, file_path):
self.data = pd.read_json(file_path, orient="records", lines=True).to_dict(orient="records")
def __len__(self):
return len(self.data)
def __getitem__(self, index):
index = min(max(index, 0), len(self.data) - 1) # 索引边界检查
# x.shape: [源序列长度]
x = torch.tensor(self.data[index]["zh_cn"], dtype=torch.long)
# y.shape: [目标序列长度]
y = torch.tensor(self.data[index]["en"], dtype=torch.long)
return x, y
创建模型：
import torch
from torch import nn
class Attention(nn.Module):
def forward(self, decoder_hidden, encoder_outputs):
# decoder_hidden.shape: [批量大小, 1, 隐藏层维度]
# encoder_outputs.shape: [批量大小, 序列长度, 隐藏层维度]
# torch.bmm(mat1, mat2) 方法对两个 3 维张量进行批量矩阵乘法运算。
# 要求 mat1 形状为 (b, n, m) 的 3 维度张量，mat2 形状为 (b, m, p) 的 3 维度张量
# 输出: 形状为 (b, n, p) 的 3 维度张量
attention_scores = torch.bmm(decoder_hidden, encoder_outputs.transpose(1, 2))
attention_widgets = torch.softmax(attention_scores, dim=-1)
return torch.bmm(attention_widgets, encoder_outputs)
class TranslationEncoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers, padding_idx):
super().__init__()
# 词嵌入层 [词向量数（词表大小）, 词向量维度, 填充值的索引]
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx)
# GRU 层 [输入维度（词向量维度）, 隐藏层维度, 隐藏层层数, 是否将batch作为输入张量的第一个维度]
self.gru = nn.GRU(embedding_dim, hidden_dim, num_layers, batch_first=True)
def forward(self, x):
# x.shape: [批量大小, 序列长度]
# embed.shape: [批量大小, 序列长度, 词向量维度]
embed = self.embedding(x) # 获取单词的嵌入向量
# output 最后一层每个时间步的隐藏层状态
# output.shape: [批量大小, 序列长度, 隐藏层维度]
# hn 最后一个时间步的隐藏层状态（包含每一个层，每一个方向）
# hn.shape: [隐藏层层数, 批量大小, 隐藏层维度]
output, hn = self.gru(embed)
lengths = (x != self.embedding.padding_idx).sum(dim=1) # 获取每个样本真实token的长度
# last_hidden_state.shape: [批量大小, 隐藏层维度]
last_hidden_state = output[torch.arange(output.shape[0]), lengths - 1] # 获取每个样本最后一个时间步隐藏状态
return output, last_hidden_state
class TranslationDecoder(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers, padding_idx):
super().__init__()
# 词嵌入层 [词表的大小, 词向量维度, 填充值的索引]
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx)
# GRU 层 [输入维度（词向量维度）, 隐藏层维度, 隐藏层层数, 是否将batch作为输入张量的第一个维度]
self.gru = nn.GRU(embedding_dim, hidden_dim, num_layers, batch_first=True)
# 添加注意力机制
self.attention = Attention()
# 输出层 [输入维度（隐藏层维度 * 2）, 输出维度（英文词表大小）]
self.liner = nn.Linear(hidden_dim * 2, vocab_size)
def forward(self, x, hx, encoder_outputs):
# x.shape: [批量大小, 序列长度（1）]
# hx.shape: [隐藏层层数（1）, 批量大小, 隐藏层维度]
# embed.shape: [批量大小, 序列长度（1）, 词向量维度]
embed = self.embedding(x) # 获取单词的嵌入向量
# output 最后一层每个时间步的隐藏层状态
# output.shape: [批量大小, 序列长度（1）, 隐藏层维度]
# hn 最后一个时间步的隐藏层状态（包含每一个层，每一个方向）
# hn.shape: [隐藏层层数, 批量大小, 隐藏层维度]
output, hn = self.gru(embed, hx)
# context_vector.shape: [批量大小, 序列长度（1）, 隐藏层维度]
context_vector = self.attention(output, encoder_outputs) # 应用注意力机制
combined = torch.cat([output, context_vector], dim=-1) # 融合信息
# Liner 层保留输入张量的所有维度（除最后一维外全部不变），仅将最后一维（特征维度）替换为 out_features
# output.shape: [批量大小, 序列长度（1）, 输出维度（英文词表大小）]
output = self.liner(combined)
return output, hn
class TranslationModel(nn.Module):
def __init__(self, vocab_sizes, embedding_dim, hidden_dim, num_layers, padding_idxs):
super().__init__()
self.encoder = TranslationEncoder(vocab_sizes[0], embedding_dim, hidden_dim, num_layers, padding_idxs[0])
self.decoder = TranslationDecoder(vocab_sizes[1], embedding_dim, hidden_dim, num_layers, padding_idxs[1])
模型单轮训练：
def model_train_one_epoch(model, device, data_loader, criterion, optimizer):
model.train() # 训练模式
total_loss = 0.0
for x, y in data_loader:
# x.shape: [批量大小, 源序列长度]
# y.shape: [批量大小, 目标序列长度]
x, y = x.to(device), y.to(device) # 指定设备
# 编码阶段
# context_vector.shape: [批量大小, 隐藏层维度]
encoder_outputs, context_vector = model.encoder(x)
# 解码阶段
decoder_inputs = y[:, : -1] # 获取解码器的输入列表
decoder_targets = y[:, 1 :] # 获取解码器的目标列表
# unsqueeze(dim) 方法在指定位置插入一个大小为 1 的新维度
# decoder_hidden.shape: [1（隐藏层层数）, 批量大小, 隐藏层维度]
decoder_hidden = context_vector.unsqueeze(0) # 获取解码器的初始状态
decoder_outputs = [] # 解码器的输出列表
seq_len = decoder_inputs.shape[1] # 获取解码器的输入序列长度
for i in range(seq_len):
# decoder_inputs[:, i].shape: [批量大小]
# decoder_input.shape: [批量大小, 1（序列长度）]
decoder_input = decoder_inputs[:, i].unsqueeze(1) # 获取当前时间步的输入
# decoder_output.shape: [批量大小, 序列长度（1）, 类别数（词汇表大小）]
decoder_output, decoder_hidden = model.decoder(decoder_input, decoder_hidden, encoder_outputs)
decoder_outputs.append(decoder_output)
# 计算损失
# decoder_targets.shape: [批量大小, 序列长度] -> [批量大小 * 序列长度]
decoder_targets = decoder_targets.reshape(-1)
# torch.cat() 就是把多个张量沿指定维度 "粘" 在一起，拼接后该维度的长度 = 各张量该维度长度之和，其它维度保持不变
# decoder_outputs: [tensor([批量大小, 序列长度（1）, 类别数（词汇表大小）])] -> [批量大小, 各个序列长度维度长度之和, 类别数（词汇表大小）]
decoder_outputs = torch.cat(decoder_outputs, dim=1) # 将多个张量按指定维度拼接
# decoder_outputs.shape：[批量大小 * 序列长度, 类别数（词汇表大小）]
decoder_outputs = decoder_outputs.reshape(-1, decoder_outputs.shape[-1])
# CrossEntropyLoss() 多分类损失函数
# 如果目标值（y）的形状为 N，则预测值（output）的形状为 [N, 类别数]
loss = criterion(decoder_outputs, decoder_targets) # 计算损失
optimizer.zero_grad() # 清空梯度
loss.backward() # 反向传播
optimizer.step() # 优化器更新参数
# x.shape[0] 是当前批次的样本数量
# 乘以 x.shape[0] 是为了消除批次大小不一致（尤其是最后一个不完整批次）带来的统计偏差
total_loss += loss.item() * x.shape[0] # 计算当前批次的损失
# len(dataloader.dataset) 返回数据集的样本数量
# len(dataloader) 返回数据加载器的批次数量
return total_loss / len(data_loader.dataset) # 返回平均损失
训练模型：
EMBEDDING_DIM = 128
HIDDEN_DIM = 256
NUM_LAYERS = 1
BATCH_SIZE = 128
LEARNING_RATE = 1e-3
EPOCHS = 10
from torch.nn.utils.rnn import pad_sequence
def collate_fn(batch):
# batch 是一个二元组列表，[(x_tensor, y_tensor)]
x_tensors = [item[0] for item in batch]
y_tensors = [item[1] for item in batch]
# pad_sequence() 方法将多个长度不同的张量（序列）填充到相同长度，并堆叠成一个批处理张量
# padding_value 填充的值，<pad> 特殊标识的索引
x_tensors = pad_sequence(x_tensors, batch_first=True, padding_value=SPECIAL_TOKENS.index(PAD_TOKEN))
y_tensors = pad_sequence(y_tensors, batch_first=True, padding_value=SPECIAL_TOKENS.index(PAD_TOKEN))
return x_tensors, y_tensors
import os
import torch
from torch.utils.data import DataLoader
def model_train(model, device, data_loader, learning_rate, epochs, pad_token_index):
model.to(device) # 指定设备
# ignore_index 让损失函数忽略指定的目标类别索引，不参与损失计算
criterion = torch.nn.CrossEntropyLoss(ignore_index=pad_token_index) # 多元交叉熵损失函数
optimizer = torch.optim.Adam(model.parameters(), learning_rate) # 创建优化器
# 开始训练
for epoch in range(epochs):
print(f"开始第 {epoch + 1} 次训练")
loss = model_train_one_epoch(model, device, data_loader, criterion, optimizer)
print(f"epoch: {epoch + 1:0>2}，loss: {loss:.4f}")
print("训练完成，开始保存模型")
if not os.path.exists(MODEL_DRIECTORY):
os.makedirs(MODEL_DRIECTORY)
file_path = os.path.join(MODEL_DRIECTORY, "translation_model_by_attention.pth")
torch.save(model.state_dict(), file_path) # 保存模型
print("模型保存完成")
if __name__ == "__main__":
pad_token_index = SPECIAL_TOKENS.index(PAD_TOKEN)
padding_idxs = [pad_token_index] * 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载词表
file_path = os.path.join(DATA_DRIECTORY, "zh_cn_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
zh_cn_vocab_list = [line.strip() for line in f.readlines()]
print("加载中文词表完成")
file_path = os.path.join(DATA_DRIECTORY, "en_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
en_vocab_list = [line.strip() for line in f.readlines()]
vocab_sizes = [len(zh_cn_vocab_list), len(en_vocab_list)]
print("加载英文词表完成")
file_path = os.path.join(DATA_DRIECTORY, "translation_train.jsonl")
train_dataset = TranslationDataset(file_path) # 创建训练数据集
print("创建训练集成功")
# collate_fn 用来指定自定义的填充方式，如果一批数据中的序列长度不一样，会报错
train_data_loader = DataLoader(train_dataset, BATCH_SIZE, True, collate_fn=collate_fn) # 创建训练数据集加载器
print("创建训练数据集加载器成功")
model = TranslationModel(vocab_sizes, EMBEDDING_DIM, HIDDEN_DIM, NUM_LAYERS, padding_idxs) # 模型
print("创建模型成功")
model_train(model, device, train_data_loader, LEARNING_RATE, EPOCHS, pad_token_index)
在机器翻译任务中，BLEU（Bilingual Evaluation Understudy） 是一种常用的自动评估指标，用于衡量模型生成的翻译与人工参考译文之间的相似程度。其核心思想是：
- n-gram 匹配：统计预测译文中有多少 n-gram（词或短语）同时出现在参考译文中，用于衡量翻译内容的准确性。
- 精确率计算：将匹配到的 n-gram 数量除以预测译文中 n-gram 的总数，反映生成译文中 “正确部分” 的比例。
此外，BLEU 还引入长度惩罚机制，防止模型通过生成过短句子获得不合理的高分。最终得到的 BLEU 分数越高，说明生成译文与参考译文越接近。
在本案例中，我们可以使用 Python 的 NLTK 库 中的 bleu_score 模块，对模型在测试集上的翻译结果进行评估，主要参考 BLEU-4 的得分情况，作为翻译质量的衡量依据。
模型单批次预测：
import torch
def model_predict_one_batch(model, device, x, en_word2index, max_seq_len):
model.eval() # 评估模式
generated_token_indexes = [] # 存储每个时间步的预测概率最大的单词索引
with torch.no_grad(): # 禁用梯度计算
# 编码阶段
# context_vector.shape: [批量大小, 隐藏层维度]
encoder_outputs, context_vector = model.encoder(x)
# 解码阶段
# decoder_hidden.shape: [1（隐藏层层数）, 批量大小, 隐藏层维度]
decoder_hidden = context_vector.unsqueeze(0) # 初始隐藏状态
batch_size = x.size(0) # 批量大小
# decoder_input.shape: [批量大小, 序列长度（1）]
# torch.full(张量, 值) 创建一个张量，所有元素都填充为指定值
decoder_input = torch.full((batch_size, 1), en_word2index[SOS_TOKEN], dtype=torch.long, device=device)
is_finished = torch.full([batch_size], False, device=device) # 记录每个样本是否生成完成
# 自回归生成
for i in range(max_seq_len):
# decoder_output.shape: [批量大小, 序列长度（1）, 词汇表大小]
decoder_output, decoder_hidden = model.decoder(decoder_input, decoder_hidden, encoder_outputs) # 解码器前向传播
# next_token_indexes.shape: [批量大小, 序列长度（1）]
next_token_indexes = decoder_output.argmax(dim=-1) # 获取每个时间步的预测概率最大的单词索引
generated_token_indexes.append(next_token_indexes) # 存储每个时间步的预测概率最大的单词索引
decoder_input = next_token_indexes # 更新解码器的输入
# 判断是否生成了 <eos> 标记
is_finished |= (next_token_indexes.squeeze(1) == en_word2index[EOS_TOKEN])
# 如果所有样本都生成完成，跳出循环
if is_finished.all():
break
# 处理预测结果
# generated_tensor.shape: [批量大小, 各个序列长度维度长度之和]
generated_tensor = torch.cat(generated_token_indexes, dim=1) # 合并所有时间步的预测概率最大的单词索引
generated_list = generated_tensor.tolist()
# 去掉 <eos> 标记之后的 token id
for index, sentence in enumerate(generated_list):
if en_word2index[EOS_TOKEN] in sentence:
generated_list[index] = sentence[:sentence.index(en_word2index[EOS_TOKEN])]
return generated_list
模型评估：
MAX_SEQ_LEN = 100
import os
import torch
from torch.utils.data import DataLoader
from nltk.translate.bleu_score import corpus_bleu
def model_evaluate(model, device, data_loader, en_word2index, max_seq_len):
model.to(device) # 指定设备
model.eval() # 评估模式
# predictions: [[*, *, *], [*, *, *]]
predictions = []
# references: [[[*, *, *]], [[*, *, *]]]
references = []
for x, y in data_loader:
# x.shape: [批量大小, 源序列长度]
# y.shape: [批量大小, 目标序列长度]
x, y = x.to(device), y.to(device) # 指定设备
# targets: [[<sos>, *, *, <eos>, <pad>], [<sos>, *, *, <eos>, <pad>], ...]
targets = y.tolist()
results = model_predict_one_batch(model, device, x, en_word2index, max_seq_len)
predictions.extend(results)
references.extend([[target[1 : target.index(en_word2index[EOS_TOKEN])]] for target in targets])
accuracy = corpus_bleu(references, predictions)
print(f"准确率：{accuracy:.4f}")
if __name__ == "__main__":
pad_token_index = SPECIAL_TOKENS.index(PAD_TOKEN)
padding_idxs = [pad_token_index] * 2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载词表
file_path = os.path.join(DATA_DRIECTORY, "zh_cn_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
zh_cn_vocab_list = [line.strip() for line in f.readlines()]
print("加载中文词表完成")
file_path = os.path.join(DATA_DRIECTORY, "en_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
en_vocab_list = [line.strip() for line in f.readlines()]
vocab_sizes = [len(zh_cn_vocab_list), len(en_vocab_list)]
print("加载英文词表完成")
# 构建词表索引
en_word2index = {word: index for index, word in enumerate(en_vocab_list)}
print("构建英文词表索引完成")
file_path = os.path.join(DATA_DRIECTORY, "translation_test.jsonl")
test_dataset = TranslationDataset(file_path) # 创建测试数据集
print("创建测试集成功")
test_data_loader = DataLoader(test_dataset, BATCH_SIZE, collate_fn=collate_fn) # 创建测试数据集加载器
print("创建测试数据集加载器成功")
model = TranslationModel(vocab_sizes, EMBEDDING_DIM, HIDDEN_DIM, NUM_LAYERS, padding_idxs).to(device)
print("创建模型成功")
file_path = os.path.join(MODEL_DRIECTORY, "translation_model_by_attention.pth")
model.load_state_dict(torch.load(file_path)) # 加载模型参数
print("加载模型参数完成")
model_evaluate(model, device, test_data_loader, en_word2index, MAX_SEQ_LEN)
模型预测：
import os
import torch
from nltk.tokenize import TreebankWordDetokenizer
def model_predict(model, device, text, zh_cn_word2index, en_word2index, en_index2word, detokenizer, max_seq_len):
model.to(device) # 指定设备
tokens = list(text) # 切词
zh_cn_indexes = [zh_cn_word2index.get(token, SPECIAL_TOKENS.index(UNK_TOKEN)) for token in tokens] # 获取单词的索引
x = torch.tensor([zh_cn_indexes], dtype=torch.long, device=device) # 输入数据
results = model_predict_one_batch(model, device, x, en_word2index, max_seq_len)
tokens = [en_index2word[index] for index in results[0]]
return detokenizer.detokenize(tokens)
def model_run_predict(model, device, word2index, en_word2index, en_index2word, max_seq_len):
print("欢迎使用中英翻译模型，输入q或quit退出")
detokenizer = TreebankWordDetokenizer()
while True:
text = input("请输入中文： \n")
if text.strip() in ["q", "quit"]:
break
if text.strip() == "":
continue
print(f"输入的中文：")
print(text)
result = model_predict(model, device, text, word2index, en_word2index, en_index2word, detokenizer, max_seq_len)
print(f"翻译结果：")
print(result)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
file_path = os.path.join(DATA_DRIECTORY, "zh_cn_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
zh_cn_vocab_list = [line.strip() for line in f.readlines()]
print("加载中文词表完成")
file_path = os.path.join(DATA_DRIECTORY, "en_vocab.txt")
with open(file_path, "r", encoding="utf-8") as f:
en_vocab_list = [line.strip() for line in f.readlines()]
vocab_sizes = [len(zh_cn_vocab_list), len(en_vocab_list)]
print("加载英文词表完成")
vocab_sizes = [len(zh_cn_vocab_list), len(en_vocab_list)]
# 构建词表到索引的映射
zh_cn_word2index = {word: index for index, word in enumerate(zh_cn_vocab_list)}
print("构建中文词表到索引的映射")
en_word2index = {word: index for index, word in enumerate(en_vocab_list)}
print("构建英文词表到索引的映射")
# 构建索引到词表的映射
en_index2word = {index: word for index, word in enumerate(en_vocab_list)}
print("构建英文词引到词表的映射")
model = TranslationModel(vocab_sizes, EMBEDDING_DIM, HIDDEN_DIM, NUM_LAYERS, padding_idxs).to(device)
print("创建模型完成")
file_path = os.path.join(MODEL_DRIECTORY, "translation_model_by_attention.pth")
model.load_state_dict(torch.load(file_path)) # 加载模型参数
print("加载模型参数完成")
model_run_predict(model, device, zh_cn_word2index, en_word2index, en_index2word, MAX_SEQ_LEN)
浙公网安备 33010602011771号