秋刀鱼在做梦

这个屌丝很懒，什么也没留下！

热门标签

基于Transformer实现机器翻译_transformer机器翻译代码

作者：秋刀鱼在做梦 | 2024-08-13 02:16:40

踩

transformer机器翻译代码

一、前言

1.1 Transformer概述

Transformer模型是一种深度学习模型，由Vaswani等人在2017年提出，主要用于自然语言处理（NLP）任务。它的核心思想是通过自注意力（Self-Attention）机制来捕捉输入数据之间的全局依赖关系，从而能够处理序列数据。

1.2 Transformer整体模型架构

二、使用Transformer和PyTorch的机器翻译模型（日译中）

2.1 配置环境以及导入软件包

确保我们的系统中安装了以下软件包，如果您发现缺少某些软件包，请务必安装它们。

由于我的电脑没有GPU，租用GPU后用于之后的模型训练。


import math
import torchtext
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from collections import Counter
from torchtext.vocab import Vocab
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io
import time
import pandas as pd
import numpy as np
import pickle
import tqdm
import sentencepiece as spm
torch.manual_seed(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(torch.cuda.get_device_name(0)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

2.2 获取平行数据集

我们将使用从JParaCrawl[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]下载的日英平行数据集，该数据集被描述为由NTT创建的最大的可公开获取的英日平行语料库。它是通过大量抓取网络并自动对齐平行句子而创建的。


df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()#[:10000]
trainja = df[3].values.tolist()#[:10000]
# trainen.pop(5972)
# trainja.pop(5972)

在导入所有日语和英语对应数据后，我删除了数据集中的最后一个数据，因为它缺少值。总的来说，trainen 和 trainja 中的句子数为 5,973,071，但是，出于学习目的，通常建议在一次性使用所有数据之前对数据进行采样并确保一切按预期工作，以节省时间。

下面是数据集中包含的句子示例：


print(trainen[500])
print(trainja[500])

2.3 准备分词器

与英语或其他字母语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的分词器，该分词器是使用SentencePiece创建的日语和英语，您可以访问JParaCrawl网站下载它们，或单击此处。


en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='spm.ja.nopretok.model')

加载分词器后，可以通过以下代码测试它们。

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type=str)

结果：

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type=str)

结果：

2.4 构建 TorchText Vocab 对象并将句子转换为 Torch 张量

使用分词器和原始句子，我们构建从 TorchText 导入的 Vocab 对象。


def build_vocab(sentences, tokenizer):#定义了一个名为 build_vocab 的函数，它接受两个参数 sentences 和 tokenizer
  counter = Counter()
  for sentence in sentences:
    counter.update(tokenizer.encode(sentence, out_type=str))#对于每个句子 sentence，它使用 tokenizer 对象的 encode 方法将句子编码成字符串类型的标记序列，并更新计数器 counter
  return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
ja_vocab = build_vocab(trainja, ja_tokenizer)
en_vocab = build_vocab(trainen, en_tokenizer)#使用 build_vocab 函数来构建日语 trainja 和英语 trainen 的词汇表，使用了对应的 ja_tokenizer 和 en_tokenizer 来对句子进行编码

获得词汇表对象后，我们可以使用词汇表和分词器对象来构建训练数据的张量。


def data_process(ja, en):
  data = []#初始化一个空列表 data，用于存储处理后的数据
  for (raw_ja, raw_en) in zip(ja, en):#使用 zip 函数迭代 ja 和 en 中对应的元素
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long)#使用 ja_tokenizer 对当前的日语句子 raw_ja 进行编码，将其转换为一个字符串标记列表
    en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    data.append((ja_tensor_, en_tensor_))#将元组 (ja_tensor_, en_tensor_) 添加到 data 列表中
  return data
train_data = data_process(trainja, trainen)#使用 data_process 函数将训练数据 (trainja 和 trainen) 处理为张量，并将结果存储在 train_data 中，用于后续的机器翻译模型训练

2.5 创建要在训练期间迭代的 DataLoader 对象


BATCH_SIZE = 8 #定义了批量处理的大小为 8，即每次从数据中取出 8 个样本进行处理
PAD_IDX = ja_vocab['<pad>'] #获取了日语词汇表中 <pad> 标记的索引，用于数据填充
BOS_IDX = ja_vocab['<bos>'] #获取了日语词汇表中 <bos>（句子开始）标记的索引
EOS_IDX = ja_vocab['<eos>'] #获取了日语词汇表中 <eos>（句子结束）标记的索引
def generate_batch(data_batch):
  ja_batch, en_batch = [], [] #初始化两个空列表 ja_batch 和 en_batch，用于存储处理后的批量数据
  for (ja_item, en_item) in data_batch:
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0)) #在日语句子张量的开头和结尾分别添加 <bos> 和 <eos> 标记，并进行张量拼接
    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0)) #在英语句子张量的开头和结尾分别添加 <bos> 和 <eos> 标记，并进行张量拼接
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX) #使用 pad_sequence 函数对日语批量 ja_batch 进行填充，使得每个批量的日语句子张量具有相同的长度，用 PAD_IDX 值进行填充
  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX) #使用 pad_sequence 函数对英语批量 en_batch 进行填充，使得每个批量的英语句子张量具有相同的长度，用 PAD_IDX 值进行填充
  return ja_batch, en_batch
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

2.6 序列到序列转换器

Transformer 是 “Attention is all you need” 论文中介绍的 Seq2Seq 模型，用于解决机器翻译任务。Transformer 模型由编码器和解码器块组成，每个块包含固定数量的层。

编码器通过一系列多头注意力和前馈网络层传播输入序列来处理输入序列。编码器的输出称为内存，与目标张量一起馈送到解码器。编码器和解码器使用教师强制技术以端到端的方式进行训练。


from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)
 
 
class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward:int = 512, dropout:float = 0.1):
        super(Seq2SeqTransformer, self).__init__() # 调用父类 nn.Module 的初始化方法
        # encoder_layer 和 decoder_layer 是 Transformer 模型的编码器层和解码器层，通过 TransformerEncoderLayer 和 TransformerDecoderLayer 初始化
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
 
        self.generator = nn.Linear(emb_size, tgt_vocab_size) #一个线性层，用于将解码器的输出转换为目标词汇的分布
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size) #输入和输出词汇表的词嵌入层
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout) #位置编码层，用于给输入的词嵌入增加位置信息
 
    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg)) #经过位置编码后的输入和输出的词嵌入
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask) #通过编码器得到的内部记忆
        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
                                        tgt_padding_mask, memory_key_padding_mask) #解码器的输出
        return self.generator(outs)
 
    #将输入序列的词嵌入经过位置编码后，通过编码器进行处理，返回编码器的输出
    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)
 
    #将目标序列的词嵌入经过位置编码后，通过解码器进行处理，返回解码器的输出
    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

文本标记通过使用标记嵌入来表示。位置编码被添加到标记嵌入中，以引入词序的概念。


class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        
        # 计算位置编码中的分母部分
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        
        # 创建位置索引张量
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        
        # 初始化位置编码张量
        pos_embedding = torch.zeros((maxlen, emb_size))
        
        # 根据位置索引和分母计算 sin 和 cos 形式的位置编码
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        
        # 扩展维度以便于与标记嵌入张量进行广播加法
        pos_embedding = pos_embedding.unsqueeze(-2)
        
        # 初始化 dropout 层和位置编码张量，并注册为模型缓冲区
        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)
 
    def forward(self, token_embedding: Tensor):
        # 在标记嵌入张量上加上位置编码，并应用 dropout
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0), :])
 
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        
        # 创建嵌入层对象
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size
    
    def forward(self, tokens: Tensor):
        # 返回经缩放后的嵌入向量
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续单词掩码来阻止目标单词关注其后续单词，同时还创建了用于屏蔽源语言和目标语言填充标记的屏蔽（mask）。


def generate_square_subsequent_mask(sz):
    # 创建一个上三角矩阵，并将其转置为下三角
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    
    # 将 mask 转换为浮点型，并将值为 0 的位置填充为负无穷，将值为 1 的位置填充为 0
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask
 
def create_mask(src, tgt):
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]
 
    # 生成目标序列的自注意力掩码
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    
    # 创建源序列的填充掩码
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)
 
    # 创建源序列的填充掩码，标记 PAD_IDX 的位置为 True
    src_padding_mask = (src == PAD_IDX).transpose(0, 1)
    
    # 创建目标序列的填充掩码，标记 PAD_IDX 的位置为 True
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
    
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

使用GPU进行训练，当使用自己的GPU的时候，NUM_ENCODER_LAYERS 和 NUM_DECODER_LAYERS 设置为3或者更高，NHEAD设置8，EMB_SIZE设置为512。


SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 512
NHEAD = 8
FFN_HID_DIM = 512
BATCH_SIZE = 16
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
NUM_EPOCHS = 16
 
# 初始化 Seq2SeqTransformer 模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)
 
# 使用 Xavier 初始化方法初始化模型参数
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)
 
# 将模型移动到指定的设备（如 GPU）
transformer = transformer.to(device)
 
# 定义损失函数为交叉熵损失函数，忽略 PAD_IDX 的预测
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)
 
# 定义优化器为 Adam 优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
 
def train_epoch(model, train_iter, optimizer):
    model.train()
    losses = 0
    for idx, (src, tgt) in enumerate(train_iter):
        src = src.to(device)
        tgt = tgt.to(device)
 
        # 生成源序列和目标输入序列的掩码
        tgt_input = tgt[:-1, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
 
        # 前向传播
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)
 
        # 梯度清零
        optimizer.zero_grad()
 
        # 生成目标输出序列
        tgt_out = tgt[1:, :]
 
        # 计算损失
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
 
        # 反向传播和优化
        loss.backward()
        optimizer.step()
 
        # 累计损失
        losses += loss.item()
 
    # 返回平均损失
    return losses / len(train_iter)
 
 
def evaluate(model, val_iter):
    model.eval()
    losses = 0
    for idx, (src, tgt) in enumerate(val_iter):
        src = src.to(device)
        tgt = tgt.to(device)
 
        # 生成源序列和目标输入序列的掩码
        tgt_input = tgt[:-1, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
 
        # 前向传播
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)
 
        # 生成目标输出序列
        tgt_out = tgt[1:, :]
 
        # 计算损失
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
 
        # 累计损失
        losses += loss.item()
 
    # 返回平均损失
    return losses / len(val_iter)

2.7 开始训练

最后，在准备了必要的类和函数之后，我们准备训练我们的模型。使用A800-80GB GPU进行训练，时间预计为55分钟。


for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
    start_time = time.time()
 
    # 训练一个 epoch
    train_loss = train_epoch(transformer, train_iter, optimizer)
 
    end_time = time.time()
 
    # 打印当前 epoch 的训练损失和训练时间
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
           f"Epoch time = {(end_time - start_time):.3f}s"))

训练过程如下：

2.8 保存词汇表对象和训练好的模型

在训练完成后，我们使用 Pickle 来保存词汇表对象（en_vocab 和 ja_vocab）。


import pickle
# open a file, where you want to store the data
file = open('en_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(en_vocab, file)
file.close()
file = open('ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()

最后，我们还可以使用 PyTorch save 和 load 函数保存模型以供以后使用。通常，有两种方法可以保存模型，具体取决于我们以后要使用它们的内容。第一个仅用于推理，我们可以稍后加载模型并使用它从日语翻译成英语。


# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

第二个参数也用于推断，但是还用于当我们想要稍后加载模型并恢复训练时。


# save model + checkpoint to resume training later
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】