文章目录
- 一、赛题链接
- 二、安装库
- 1.spacy
- 2.torch_text
- 三、数据预处理
- 赛题数据
- 类定义 `TranslationDataset`
- 批量处理函数 `collate_fn`
- 四、编码器和解码器
- Encoder 类
- Decoder 类
- Seq2Seq 类
- 注意事项
- 五、主函数
- 1. `load_terminology_dictionary(dict_file)`
- 2. `train(model, iterator, optimizer, criterion, clip)`
- 六、验证集的模型评价和测试集的模型推理
一、赛题链接
赛题链接:https://challenge.xfyun.cn/topic/info?type=machine-translation-2024
二、安装库
1.spacy
查看本地spacy版本
pip show spacy
我安装3.6.0
pip install en_core_web_sm-3.6.0.tar.gz
en_core_web_sm下载链接:https://github.com/explosion/spacy-models/releases
2.torch_text
!pip install torchtext
命令 !pip install torchtext 是一个在支持Jupyter Notebook或类似环境的Python解释器中使用的命令,用于安装或更新torchtext库。这个命令通过Python的包管理工具pip来执行。
!:这个符号在Jupyter Notebook、Google Colab等环境中用作前缀,允许你在代码单元中执行shell命令。这意呀着,紧跟在这个符号后面的命令将会作为shell命令执行,而不是Python代码。
torchtext是PyTorch生态系统中的一个库,它提供了一套用于处理自然语言和其他文本数据的工具,包括数据加载、预处理、词汇表构建和简单的文本分类等。
三、数据预处理
赛题数据
- 训练集:双语数据 - 中英14万余双语句对
- 开发集:英中1000双语句对
- 测试集:英中1000双语句对
- 术语词典:英中2226条
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchtext.data.utils import get_tokenizer
from collections import Counter
import random
from torch.utils.data import Subset, DataLoader
import time
# 定义数据集类
# 修改TranslationDataset类以处理术语
class TranslationDataset(Dataset):
def __init__(self, filename, terminology):
self.data = []
with open(filename, 'r', encoding='utf-8') as f:
for line in f:
en, zh = line.strip().split('\t')
self.data.append((en, zh))
self.terminology = terminology
# 创建词汇表,注意这里需要确保术语词典中的词也被包含在词汇表中
self.en_tokenizer = get_tokenizer('basic_english')
self.zh_tokenizer = list # 使用字符级分词
en_vocab = Counter(self.terminology.keys()) # 确保术语在词汇表中
zh_vocab = Counter()
for en, zh in self.data:
en_vocab.update(self.en_tokenizer(en))
zh_vocab.update(self.zh_tokenizer(zh))
# 添加术语到词汇表
self.en_vocab = ['<pad>', '<sos>', '<eos>'] + list(self.terminology.keys()) + [word for word, _ in en_vocab.most_common(10000)]
self.zh_vocab = ['<pad>', '<sos>', '<eos>'] + [word for word, _ in zh_vocab.most_common(10000)]
self.en_word2idx = {word: idx for idx, word in enumerate(self.en_vocab)}
self.zh_word2idx = {word: idx for idx, word in enumerate(self.zh_vocab)}
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
en, zh = self.data[idx]
en_tensor = torch.tensor([self.en_word2idx.get(word, self.en_word2idx['<sos>']) for word in self.en_tokenizer(en)] + [self.en_word2idx['<eos>']])
zh_tensor = torch.tensor([self.zh_word2idx.get(word, self.zh_word2idx['<sos>']) for word in self.zh_tokenizer(zh)] + [self.zh_word2idx['<eos>']])
return en_tensor, zh_tensor
def collate_fn(batch):
en_batch, zh_batch = [], []
for en_item, zh_item in batch:
en_batch.append(en_item)
zh_batch.append(zh_item)
# 对英文和中文序列分别进行填充
en_batch = nn.utils.rnn.pad_sequence(en_batch, padding_value=0, batch_first=True)
zh_batch = nn.utils.rnn.pad_sequence(zh_batch, padding_value=0, batch_first=True)
return en_batch, zh_batch
这段代码定义了一个用于机器翻译任务的数据集类 TranslationDataset
,它继承自 PyTorch 的 Dataset
类。这个类特别设计来处理包含英文和中文翻译对的文本文件,并且它还支持一个术语词典(terminology
),以确保这些术语在构建词汇表时被优先考虑。下
类定义 TranslationDataset
-
初始化方法
__init__
:- 读取翻译文件(每行包含一个英文句子和一个中文句子,由制表符分隔)。
- 初始化术语词典(
terminology
),这是一个字典。 - 使用
get_tokenizer('basic_english')
获取英文的基本分词器。对于中文,这里简单地将整个句子视为一个字符序列(使用list
作为分词器,实际上并不是真正的分词,但在这个示例中为了简化处理)。 - 使用
Counter
计数英文术语、英文句子分词结果和中文句子分词结果(尽管中文这里按字符处理,但仍进行计数以便后续筛选高频词)。 - 构建英文和中文的词汇表。首先添加特殊标记(
<pad>
,<sos>
,<eos>
),然后添加术语词典中的词,最后添加最常见的10000个英文和中文词(通过most_common(10000)
获取)。 - 创建词汇表到索引的映射(
word2idx
)。
-
长度方法
__len__
:- 返回数据集中翻译对的数量。
-
获取项方法
__getitem__
:- 根据索引
idx
获取数据集中的一个翻译对(英文和中文)。 - 使用分词器将英文句子分词,并将每个词转换为索引。
- 对于中文句子,由于这里按字符处理,直接将每个字符转换为索引。
- 序列开始和结束标记:在序列两端添加
<sos>
(Sequence Start)和<eos>
(Sequence End)标记,帮助模型识别序列的起始和结束。 - 在英文和中文句子的末尾添加
<eos>
索引表示句子结束。 - 返回英文和中文句子的索引张量。
- 根据索引
批量处理函数 collate_fn
- 这个函数用于在
DataLoader
中将多个样本组合成一个批次。 - 它遍历批次中的每个英文和中文句子索引张量,并将它们分别收集到
en_batch
和zh_batch
中。 - 使用
nn.utils.rnn.pad_sequence
对英文和中文句子批次进行填充,以确保批次中的每个句子都有相同的长度(较短的句子用0填充,即<pad>
的索引)。这里设置batch_first=True
表示批次维度是第一维。 - 返回填充后的英文和中文句子批次。
四、编码器和解码器
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
super().__init__()
self.embedding = nn.Embedding(input_dim, emb_dim)
self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
# src shape: [batch_size, src_len]
embedded = self.dropout(self.embedding(src))
# embedded shape: [batch_size, src_len, emb_dim]
outputs, hidden = self.rnn(embedded)
# outputs shape: [batch_size, src_len, hid_dim]
# hidden shape: [n_layers, batch_size, hid_dim]
return outputs, hidden
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
super().__init__()
self.output_dim = output_dim
self.embedding = nn.Embedding(output_dim, emb_dim)
self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
self.fc_out = nn.Linear(hid_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, input, hidden):
# input shape: [batch_size, 1]
# hidden shape: [n_layers, batch_size, hid_dim]
embedded = self.dropout(self.embedding(input))
# embedded shape: [batch_size, 1, emb_dim]
output, hidden = self.rnn(embedded, hidden)
# output shape: [batch_size, 1, hid_dim]
# hidden shape: [n_layers, batch_size, hid_dim]
prediction = self.fc_out(output.squeeze(1))
# prediction shape: [batch_size, output_dim]
return prediction, hidden
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, trg, teacher_forcing_ratio=0.5):
# src shape: [batch_size, src_len]
# trg shape: [batch_size, trg_len]
batch_size = src.shape[0]
trg_len = trg.shape[1]
trg_vocab_size = self.decoder.output_dim
outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)
_, hidden = self.encoder(src)
input = trg[:, 0].unsqueeze(1) # Start token
for t in range(1, trg_len):
output, hidden = self.decoder(input, hidden)
outputs[:, t, :] = output
teacher_force = random.random() < teacher_forcing_ratio
top1 = output.argmax(1)
input = trg[:, t].unsqueeze(1) if teacher_force else top1.unsqueeze(1)
return outputs
上述代码实现了基于GRU(门控循环单元)的序列到序列(Seq2Seq)模型,该模型通常用于机器翻译、文本摘要等任务。代码由三个主要部分组成:Encoder
类、Decoder
类和 Seq2Seq
类。下面是对每个部分的详细解释:
Encoder 类
- 初始化 (
__init__
): 接收输入维度(input_dim
)、嵌入维度(emb_dim
)、隐藏层维度(hid_dim
)、层数(n_layers
)和dropout比率(dropout
)作为参数。构建了一个嵌入层(nn.Embedding
)用于将输入转换为嵌入向量,一个GRU层(nn.GRU
)用于处理序列数据,以及一个dropout层(nn.Dropout
)用于减少过拟合。 - 前向传播 (
forward
): 接收源序列(src
)作为输入,首先通过嵌入层和dropout层,然后通过GRU层。返回GRU的输出和最终的隐藏状态。
Decoder 类
- 初始化 (
__init__
): 类似于Encoder,但增加了一个全连接层(nn.Linear
)用于将GRU的输出转换为预测的输出维度。 - 前向传播 (
forward
): 接收目标序列的当前输入(input
)和隐藏状态(hidden
)作为输入。首先将输入通过嵌入层和dropout层,然后通过GRU层。最后,使用全连接层将GRU的输出转换为预测,并返回预测和更新后的隐藏状态。
Seq2Seq 类
- 初始化 (
__init__
): 接收Encoder和Decoder实例以及设备(device
)作为参数。这些实例和设备将用于模型的训练和推理。 - 前向传播 (
forward
): 接收源序列(src
)和目标序列(trg
)作为输入,以及一个可选的教师强制比率(teacher_forcing_ratio
)。首先,通过Encoder处理源序列以获取初始隐藏状态。然后,使用目标序列的第一个元素(通常是开始标记)作为Decoder的初始输入。在接下来的时间步中,Decoder根据当前输入和隐藏状态生成输出。根据教师强制比率,Decoder的下一个输入可以是目标序列的实际下一个元素(教师强制)或当前时间步的预测(非教师强制)。最后,返回所有时间步的预测输出。
注意事项
teacher_forcing_ratio
用于在训练过程中平衡教师强制和非教师强制的比例。教师强制有助于模型更快地学习,因为它总是提供正确的下一个输入;然而,非教师强制有助于模型学习在测试时如何自己生成序列。- 模型的输出
outputs
是一个三维张量,其中包含了每个时间步每个样本在每个词汇上的预测概率。
五、主函数
# 新增术语词典加载部分
def load_terminology_dictionary(dict_file):
terminology = {}
with open(dict_file, 'r', encoding='utf-8') as f:
for line in f:
en_term, ch_term = line.strip().split('\t')
terminology[en_term] = ch_term
return terminology
def train(model, iterator, optimizer, criterion, clip):
model.train()
epoch_loss = 0
for i, (src, trg) in enumerate(iterator):
src, trg = src.to(device), trg.to(device)
optimizer.zero_grad()
output = model(src, trg)
output_dim = output.shape[-1]
output = output[:, 1:].contiguous().view(-1, output_dim)
trg = trg[:, 1:].contiguous().view(-1)
loss = criterion(output, trg)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(iterator)
这段代码包含两个主要的函数,用于处理术语词典的加载和模型训练的流程。下面是对这两个函数的详细解释:
1. load_terminology_dictionary(dict_file)
这个函数用于加载一个术语词典文件,并将文件中的内容转换为一个Python字典。这个字典的键(key)是英文术语,值(value)是对应的中文术语。
-
参数:
dict_file
: 术语词典文件的路径,该文件每行包含一个英文术语和一个中文术语,两者之间用制表符(\t
)分隔。
-
过程:
- 初始化一个空字典
terminology
来存储术语对。 - 使用
with open(...)
语句以只读模式('r'
)和UTF-8编码打开术语词典文件。这样做可以确保文件正确关闭,即使在读取文件时发生异常也是如此。 - 遍历文件的每一行,使用
strip()
方法去除每行末尾的换行符等空白字符,然后使用split('\t')
方法将每行按制表符分割成英文术语和中文术语。 - 将英文术语作为键,中文术语作为值,存入
terminology
字典中。 - 遍历完成后,返回
terminology
字典。
- 初始化一个空字典
2. train(model, iterator, optimizer, criterion, clip)
这个函数定义了模型训练的一个完整周期(epoch)的流程。它接收一个模型、一个数据迭代器、一个优化器、一个损失函数和一个梯度裁剪值作为参数。
-
参数:
model
: 待训练的模型。iterator
: 数据迭代器,用于遍历训练数据。optimizer
: 优化器,用于更新模型的参数以最小化损失函数。criterion
: 损失函数,用于评估模型预测和真实标签之间的差异。clip
: 梯度裁剪的阈值,用于防止梯度爆炸。
-
过程:
- 将模型设置为训练模式(
model.train()
)。 - 初始化
epoch_loss
为0,用于记录整个训练周期的总损失。 - 遍历迭代器中的每一批数据(
src
,trg
),其中src
是源语言数据,trg
是目标语言数据。 - 将源数据和目标数据移动到指定的设备(如GPU)上。
- 清零优化器的梯度。
- 通过模型进行前向传播,得到预测结果
output
。 - 由于模型通常输出的是整个序列的预测(包括起始标记),而损失计算通常不包括起始标记,因此需要调整
output
和trg
的形状,以排除起始标记。 - 计算损失值
loss
。 - 通过反向传播计算梯度。
- 使用梯度裁剪来防止梯度爆炸。
- 更新模型的参数。
- 累加当前批次的损失值到
epoch_loss
。 - 遍历完成后,计算并返回整个训练周期的平均损失值。
- 将模型设置为训练模式(
# 主函数
if __name__ == '__main__':
start_time = time.time() # 开始计时
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#terminology = load_terminology_dictionary('../dataset/en-zh.dic')
terminology = load_terminology_dictionary('../dataset/en-zh.dic')
# 加载数据
dataset = TranslationDataset('../dataset/train.txt',terminology = terminology)
# 选择数据集的前N个样本进行训练
N = 1000 #int(len(dataset) * 1) # 或者你可以设置为数据集大小的一定比例,如 int(len(dataset) * 0.1)
subset_indices = list(range(N))
subset_dataset = Subset(dataset, subset_indices)
train_loader = DataLoader(subset_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
# 定义模型参数
INPUT_DIM = len(dataset.en_vocab)
OUTPUT_DIM = len(dataset.zh_vocab)
ENC_EMB_DIM = 256
DEC_EMB_DIM = 256
HID_DIM = 512
N_LAYERS = 2
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5
# 初始化模型
enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)
model = Seq2Seq(enc, dec, device).to(device)
# 定义优化器和损失函数
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=dataset.zh_word2idx['<pad>'])
# 训练模型
N_EPOCHS = 10
CLIP = 1
for epoch in range(N_EPOCHS):
train_loss = train(model, train_loader, optimizer, criterion, CLIP)
print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')
# 在训练循环结束后保存模型
torch.save(model.state_dict(), './translation_model_GRU.pth')
end_time = time.time() # 结束计时
# 计算并打印运行时间
elapsed_time_minute = (end_time - start_time)/60
print(f"Total running time: {elapsed_time_minute:.2f} minutes")
下面是对代码主要部分的详细解释:
-
环境设置与设备选择:
- 使用
torch.device
来检查CUDA是否可用,并据此选择使用GPU还是CPU进行模型训练。
- 使用
-
术语词典加载:
- 调用
load_terminology_dictionary
函数加载一个术语词典文件(如en-zh.dic
),该文件包含英文术语及其对应的中文翻译。加载了术语词典到terminology
变量中。
- 调用
-
数据加载与预处理:
- 使用
TranslationDataset
类(加载训练数据集(如train.txt
),并传入术语词典。 - 从数据集中选择前
N
个样本(这里是1000个)进行训练,通过Subset
和DataLoader
进行批量处理和打乱数据。
- 使用
-
模型参数定义:
- 根据数据集的词汇表大小等设置模型的输入维度、输出维度、嵌入维度、隐藏层维度、层数及dropout率等参数。
-
模型初始化:
- 分别初始化编码器(
Encoder
)和解码器(Decoder
),然后将它们组合成Seq2Seq
模型,并将模型移至选定的设备上(GPU或CPU)。
- 分别初始化编码器(
-
优化器和损失函数:
- 使用Adam优化器来优化模型参数。
- 使用交叉熵损失函数(忽略填充索引)来计算预测和真实标签之间的差异。
-
模型训练:
- 通过一个训练循环,进行多个epoch的训练。在每个epoch中,调用
train
函数来训练模型。 - 使用梯度裁剪(
CLIP
)来防止梯度爆炸。 - 打印每个epoch的训练损失。
- 通过一个训练循环,进行多个epoch的训练。在每个epoch中,调用
-
模型保存:
- 在所有epoch完成后,保存模型的参数到文件中(如
translation_model_GRU.pth
)。
- 在所有epoch完成后,保存模型的参数到文件中(如
-
运行时间计算:
- 计算并打印从程序开始到结束的总运行时间(以分钟为单位)。
六、验证集的模型评价和测试集的模型推理
见教程
https://datawhaler.feishu.cn/wiki/FVs2wAVN5iqHMqk5lW2ckfhAncb