如何利用PyTorch写一个Transformer实现英德互译
数据集中每一行是一对英语，德语句子对
Transformer模型出处：2017 《Attention is all you need》
Transformer中的位置编码是什么意思？
https://kazemnejad.com/blog/transformer_architecture_positional_encoding/kazemnejad.com
运行以下代码请确保：
PyTorch=1.9.0
torchtext=0.10.0
否则肯定报错
# 利用Transformer结构
# 实现英语-德语互译
# 分词工具
from torchtext.data.utils import get_tokenizer
# 构建数据集的词典的工具
from torchtext.vocab import build_vocab_from_iterator
# Multi30k数据集
# 常用的机器翻译数据集
from torchtext.datasets import Multi30k
from typing import Iterable, List
# 从德语翻译到英语
SRC_LANGUAGE = 'de'
TGT_LANGUAGE = 'en'
# 转换器
token_transform = {}
vocab_transform = {}
# 请确保安装了以下依赖包
# pip install -U spacy
# 定义指定的分词格式
# 按照以下网址中的安装方法进行安装
# 下载指定的分词格式
# https://blog.csdn.net/yyp1998/article/details/117292281
# get_torkenizer(第一个参数指定使用的分词器，如果没有,按照空格进行分割，language=指定使用的分词格式（哪种语言）)
# 德语分词器
token_transform[SRC_LANGUAGE] = get_tokenizer('spacy', language='de_core_news_sm')
# 英语分词器
token_transform[TGT_LANGUAGE] = get_tokenizer('spacy', language='en_core_web_sm')
def yield_tokens(data_iter: Iterable, language: str) -> List[str]:
# 利用上面的分词器
# 对数据集中的每行英语<->德语对进行分词处理
# 标记是英语还是德语
language_index = {SRC_LANGUAGE: 0, TGT_LANGUAGE: 1}
# 迭代数据集中的每一行
for data_sample in data_iter:
# 产生这一行的分词
yield token_transform[language](data_sample[language_index[language]])
# 定义特殊标记
# UNK=unknown 未知词语标记
# PAD=padding 序列填充标记
# BOS=begin of string(有的也定义为SOS，start of string) 一句话的开始标记
# EOS=end of string 一句话的结尾标记
UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3
special_symbols = ['', '', '', '']
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
# split指定只获取训练数据集
# language_pair指定我们想要使用的句子对
train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
# 根据数据集建立一个词典
# 加入我们上面定义的特殊标记
# 并把这些特殊标记放在词典的最前面
# 最小词频指定为1，意思是只要训练集中出现过的词语，通通加入到词典中
vocab_transform[ln] = build_vocab_from_iterator(yield_tokens(train_iter, ln),
min_freq=1,
specials=special_symbols,
special_first=True)
# 配置UNK_IDX为默认索引
# 后面翻译过程中要是出现词典中没见过的词语，一律判定为UNK
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
vocab_transform[ln].set_default_index(UNK_IDX)
# 测试
# 获取am begin在我们建立的英语词典中的索引
print(vocab_transform['en'](['am', 'begin']))
from torch import Tensor
import torch
import torch.nn as nn
from torch.nn import Transformer
import math
# 获取可用设备
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 位置编码
# 详细的看我的上几篇中的《如何利用PyTorch训练一个Transformer语言模型学习词嵌入》
class PositionalEncoding(nn.Module):
def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000):
'''
:param emb_size: 词嵌入的维度
:param dropout: dropout概率
:param maxlen: 一句话的最大长度
'''
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))
pos_embedding[:, 0::2] = torch.sin(pos * den)
pos_embedding[:, 1::2] = torch.cos(pos * den)
# 添加一个批大小维度
pos_embedding = pos_embedding.unsqueeze(-2)
self.dropout = nn.Dropout(dropout)
self.register_buffer('pos_embedding', pos_embedding)
def forward(self, token_embedding: Tensor):
return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])
# 词嵌入
class TokenEmbedding(nn.Module):
def __init__(self, vocab_size: int, emb_size):
'''
:param vocab_size: 词典大小
:param 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)
# Transformer
class Seq2SeqTransformer(nn.Module):
def __init__(self, num_encoder_layers: int, num_decoder_layers: int, emb_size: int, nhead: int, src_vocab_size: int,
tgt_vocab_size: int, dim_feedforward: int = 512, dropout: float = 0.1):
'''
:param num_encoder_layers: 编码器层数
:param num_decoder_layers: 解码器层数
:param emb_size: 词嵌入维度
:param nhead: 多头注意力头数
:param src_vocab_size: 源语言词典大小
:param tgt_vocab_size: 目标语言词典大小
:param dim_feedforward: 隐藏层的维度大小
:param dropout: dropout概率
'''
super(Seq2SeqTransformer, self).__init__()
self.transformer = Transformer(d_model=emb_size,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward,
dropout=dropout)
# 输出层
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):
'''
:param src: 输入的句子
:param trg: 想要翻译成的句子
:param src_mask: 对于输入的Mask矩阵
:param tgt_mask: 对于输出的Mask矩阵
:param src_padding_mask:对于填充的Mask矩阵
:param tgt_padding_mask:
:param memory_key_padding_mask:Encoder输出的Mask矩阵
:return:
'''
# 先进行词嵌入，然后进行位置编码
src_emb = self.positional_encoding(self.src_tok_emb(src))
tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
# 输入到transformer中进行编、解码
outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
# 产生transformer的输出
return self.generator(outs)
# Encoder
def encode(self, src: Tensor, src_mask: Tensor):
return self.transformer.encoder(self.positional_encoding(
self.src_tok_emb(src)), src_mask)
# Decoder
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)
def generate_square_subsequent_mask(sz):
# torch.ones((sz, sz), device=DEVICE)生成一个指定设备上元素都为1的[sz,sz]矩阵
# torch.triu()函数用于从矩阵中沿着主对角线截取一部分元素，diagonal=参数控制截取的位置
# 如果不指定默认diagonal=0,就是截取一个矩阵的上三角部分
# torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1转换为一个bool值的mask矩阵
# .transpose(0,1)是矩阵转置
mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)
# 将False元素设置为-inf
# 其余元素设置为0.0
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
# 上面的内容可以写为
# mask=torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)
return mask
print(generate_square_subsequent_mask(3))
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)
# 如果指定了Padding，进行全0填充
# 词嵌入中的padding是做什么的呢
# 为了方便模型的训练
# 如果某句话的长度没有达到这批次数据每句话的最大长度，进行padding填充对齐
src_padding_mask = (src == PAD_IDX).transpose(0, 1)
tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask
torch.manual_seed(0)
# 词典的大小
SRC_VOCAB_SIZE = len(vocab_transform[SRC_LANGUAGE])
TGT_VOCAB_SIZE = len(vocab_transform[TGT_LANGUAGE])
# 词嵌入维度
EMB_SIZE = 512
# 多头注意力头数
NHEAD = 8
# Transformer隐藏层的维度
FFN_HID_DIM = 512
# 批大小
BATCH_SIZE = 128
# 编、解码器的层数
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
# 实例化一个Transformer
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
# 权重初始化
for p in transformer.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
# 如果GPU可用放在GPU上
transformer = transformer.to(DEVICE)
# 损失函数
# padding的不是真正的词语，要忽略掉 ignore_index=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)
from torch.nn.utils.rnn import pad_sequence
# 转换器聚合器
# 组合多个转换器对输入进行转换
# 类似torchvision中的Compose
# 它的原理是什么呢
# 首先*transforms这个动作表示解包操作，拆分输入进来的多个转换器
# 然后通过内置函数依次利用不同的转换器处理输入
# 相当于一个递归，但是每次调用func时使用的转换器不同
def sequential_transforms(*transforms):
def func(txt_input):
for transform in transforms:
txt_input = transform(txt_input)
return txt_input
return func
# 在一句话的开头和结尾添加特殊标记
def tensor_transform(token_ids: List[int]):
return torch.cat((torch.tensor([BOS_IDX]),
torch.tensor(token_ids),
torch.tensor([EOS_IDX])))
# 组合转换器
text_transform = {}
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
# 先分词
# 然后获取每个词语在词典中的索引
# 然后转换为张量
text_transform[ln] = sequential_transforms(token_transform[ln],
vocab_transform[ln],
tensor_transform)
# 处理一批次的数据
def collate_fn(batch):
src_batch, tgt_batch = [], []
for src_sample, tgt_sample in batch:
# rstrip()删除字符串末尾的指定字符
# 这两行在做什么呢
# 根据我们前面组合的转换器
# 对于一句话
# 先分词
# 然后获取这句话每个词在词典中的索引
# 然后加入开始，结束标记，转换为张量
# 然后加入到这批次数据中
src_batch.append(text_transform[SRC_LANGUAGE](src_sample.rstrip("\n")))
tgt_batch.append(text_transform[TGT_LANGUAGE](tgt_sample.rstrip("\n")))
# 拿src_batch举例
# 处理完后src_batch是一个长度为batch_size的列表，其中每个元素是[seq_len,源语言词典大小]的矩阵
# 但是注意这里seq_len长度不一定都相同
# 因为你无法保证每一句话的长度都是相同的
# 所以利用pad_sequence进行填充，填充到相同的长度
# 方便处理
# 但是这些填充又不是真正的词语
# 所以我们要引入Padding Mask矩阵遮盖住这些地方，防止对翻译产生影响
src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)
tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)
# pad_sequence之后
# src_batch的维度为[统一的seq_len,batch_size,源语言词典大小]
return src_batch, tgt_batch
from torch.utils.data import DataLoader
def train_epoch(model, optimizer):
# 开启模型训练模式
model.train()
losses = 0
# 获取一批次的数据
# 并指定对于这批数据使用我们上面定义的处理方法， 转换为张量
train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
train_dataloader = DataLoader(train_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)
# 训练
for src, tgt in train_dataloader:
# 如果GPU可用放在GPU上
src = src.to(DEVICE)
tgt = tgt.to(DEVICE)
# 最后一个是结尾特殊标志
# 省去
tgt_input = tgt[:-1, :]
# 获取Mask矩阵
src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
# 获取模型的输出
# logits是没有经过softmax的模型输出的意思
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:, :]
# logits的维度是[seq_len-1,batch_size,目标语言的词典大小]
# logits.reshape(-1, logits.shape[-1])后维度变为[(seq_len-1)*batch_size,目标语言的词典大小]
# tgt_out维度是[seq_len-1,batch_size,目标语言的词典大小]
# tgt_out.reshape(-1)后tgt_out的维度是[seq_len-1*batch_size*目标语言的词典大小]
# 为什么要这么做呢
# 可以类比图片分类
# tgt_out的最后一个维度是真实标签
# 而logits的最后一个维度是对于每一个词语的预测
loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
# 反向传播
loss.backward()
# 梯度更新
optimizer.step()
losses += loss.item()
return losses / len(train_dataloader)
# 模型验证
def evaluate(model):
model.eval()
losses = 0
val_iter = Multi30k(split='valid', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
val_dataloader = DataLoader(val_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)
for src, tgt in val_dataloader:
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_dataloader)
from timeit import default_timer as timer
NUM_EPOCHS = 18
# 迭代训练和验证
for epoch in range(1, NUM_EPOCHS + 1):
start_time = timer()
train_loss = train_epoch(transformer, optimizer)
end_time = timer()
val_loss = evaluate(transformer)
print((
f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, "f"Epoch time = {(end_time - start_time):.3f}s"))
# 让训练完的模型产生翻译字符串
def greedy_decode(model, src, src_mask, max_len, start_symbol):
src = src.to(DEVICE)
src_mask = src_mask.to(DEVICE)
memory = model.encode(src, src_mask)
ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(DEVICE)
for i in range(max_len - 1):
memory = memory.to(DEVICE)
tgt_mask = (generate_square_subsequent_mask(ys.size(0))
.type(torch.bool)).to(DEVICE)
out = model.decode(ys, memory, tgt_mask)
out = out.transpose(0, 1)
prob = model.generator(out[:, -1])
_, next_word = torch.max(prob, dim=1)
next_word = next_word.item()
ys = torch.cat([ys,
torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
if next_word == EOS_IDX:
break
return ys
# 翻译
def translate(model: torch.nn.Module, src_sentence: str):
model.eval()
src = text_transform[SRC_LANGUAGE](src_sentence).view(-1, 1)
num_tokens = src.shape[0]
src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
tgt_tokens = greedy_decode(
model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
return " ".join(vocab_transform[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("",
"").replace(
"", "")
# 将德语翻译为英语
print(translate(transformer, "Eine Gruppe von Menschen steht vor einem Iglu ."))