本文将会逐一拆解Transformer网络，并给出代码实现，最终将应用Transformer模型完成一个文本翻译任务。
Transformer网络拆解
Transformer唯一难理解的地方，也是唯一涉及到了数学知识的地方就是正弦位置编码表，对于这部分的工作原理大可不比过于纠结，知道它是做什么的、会用就可以；其他地方只需要像搭积木一样摆放到合适的位置就好。
缩放点积注意力
虽然说Attention机制是Transformer的核心，但是这个机制并不复杂，无非就是两次矩阵相乘
d_k和d_v变量表示的是分头后，每个头的向量维度，目前设置的值64为Embedding的维度256除以头数量8
d_k = 64 # K(=Q)维度
d_v = 64 # V维度
# 定义缩放点积注意力类
class ScaledDotProductAttention(nn.Module):
def __init__(self):
super(ScaledDotProductAttention, self).__init__()
def forward(self, Q, K, V, attn_mask):
#-------------------------维度信息--------------------------------
# Q K V [batch_size, n_heads, len_q/k/v, dim_q=k/v] (dim_q=dim_k)
# attn_mask [batch_size, n_heads, len_q, len_k]
#-----------------------------------------------------------------
# 计算注意力分数（原始权重）[batch_size，n_heads，len_q，len_k]
scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)
#-------------------------维度信息--------------------------------
# scores [batch_size, n_heads, len_q, len_k]
#-----------------------------------------------------------------
# 使用注意力掩码，将attn_mask中值为1的位置的权重替换为极小值
#-------------------------维度信息--------------------------------
# attn_mask [batch_size, n_heads, len_q, len_k], 形状和scores相同
#-------------------------维度信息--------------------------------
scores.masked_fill_(attn_mask, -1e9)
# 对注意力分数进行softmax
weights = nn.Softmax(dim=-1)(scores)
#-------------------------维度信息--------------------------------
# weights [batch_size, n_heads, len_q, len_k], 形状和scores相同
#-------------------------维度信息--------------------------------
# 计算上下文向量（也就是注意力的输出）, 是上下文信息的紧凑表示
context = torch.matmul(weights, V)
#-------------------------维度信息--------------------------------
# context [batch_size, n_heads, len_q, dim_v]
#-------------------------维度信息--------------------------------
return context, weights # 返回上下文向量和注意力分数多头自注意力
# 定义多头注意力类
d_embedding = 512 # Embedding的维度
n_heads = 8 # Multi-Head Attention中头的个数
class MultiHeadAttention(nn.Module):
def __init__(self):
super(MultiHeadAttention, self).__init__()
self.W_Q = nn.Linear(d_embedding, d_k * n_heads) # Q的线性变换层
self.W_K = nn.Linear(d_embedding, d_k * n_heads) # K的线性变换层
self.W_V = nn.Linear(d_embedding, d_v * n_heads) # V的线性变换层
self.linear = nn.Linear(n_heads * d_v, d_embedding)
self.layer_norm = nn.LayerNorm(d_embedding)
def forward(self, Q, K, V, attn_mask):
#-------------------------维度信息--------------------------------
# Q K V [batch_size, len_q/k/v, embedding_dim]
#-------------------------维度信息--------------------------------
residual, batch_size = Q, Q.size(0) # 保留残差连接
# 将输入进行线性变换和重塑，以便后续处理
q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)
k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2)
v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)
#-------------------------维度信息--------------------------------
# q_s k_s v_s: [batch_size, n_heads, len_q/k/v, d_q=k/v]
#-------------------------维度信息--------------------------------
# 将注意力掩码复制到多头 attn_mask: [batch_size, n_heads, len_q, len_k]
attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)
#-------------------------维度信息--------------------------------
# attn_mask [batch_size, n_heads, len_q, len_k]
#-------------------------维度信息--------------------------------
# 使用缩放点积注意力计算上下文和注意力权重
# context: [batch_size, n_heads, len_q, dim_v]; weights: [batch_size, n_heads, len_q, len_k]
context, weights = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
#-------------------------维度信息--------------------------------
# context [batch_size, n_heads, len_q, dim_v]
# weights [batch_size, n_heads, len_q, len_k]
#-------------------------维度信息--------------------------------
# 重塑上下文向量并进行线性变换，[batch_size, len_q, n_heads * dim_v]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)
#-------------------------维度信息--------------------------------
# context [batch_size, len_q, n_heads * dim_v]
#-------------------------维度信息--------------------------------
output = self.linear(context) # [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# output [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# 与输入(Q)进行残差链接，并进行层归一化后输出[batch_size, len_q, embedding_dim]
output = self.layer_norm(output + residual)
#-------------------------维度信息--------------------------------
# output [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
return output, weights # 返回层归一化的输出和注意力权重前馈神经网络
这里使用了两个一维卷积对数据进行了升维和降维，在进行卷积操作前需要进行一个转置操作，原因是：一维卷积其实可以看成是一种特殊的二维卷积，特殊体现在卷积核的长度与embedding_dim相同，而宽度则是kernel_size的数值，这样的卷积核只能在一个方向上滑动，所以被称为一维卷积，滑动的方向为张量的最后一个维度；在NLP任务中我们希望卷积核沿着句子的方向一个词一个词地移动，进行转置操作可以将seq_length放到最后。
# 定义逐位置前向传播网络类
class PoswiseFeedForwardNet(nn.Module):
def __init__(self):
super(PoswiseFeedForwardNet, self).__init__()
# 定义一维卷积层1，用于将输入映射到更高维度
self.conv1 = nn.Conv1d(in_channels=d_embedding, out_channels=2048, kernel_size=1)
# 定义一维卷积层2，用于将输入映射回原始维度
self.conv2 = nn.Conv1d(in_channels=2048, out_channels=d_embedding, kernel_size=1)
# 定义层归一化
self.layer_norm = nn.LayerNorm(d_embedding)
def forward(self, inputs): # inputs: [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# inputs [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
residual = inputs # 保留残差连接 [batch_size, len_q, embedding_dim]
# 在卷积层1后使用ReLU激活函数 [batch_size, embedding_dim, len_q]->[batch_size, 2048, len_q]
output = nn.ReLU()(self.conv1(inputs.transpose(1, 2)))
#-------------------------维度信息--------------------------------
# output [batch_size, 2048, len_q]
#-------------------------维度信息--------------------------------
# 使用卷积层2进行降维 [batch_size, 2048, len_q]->[batch_size, embedding_dim, len_q]
output = self.conv2(output).transpose(1, 2) # [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# output [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# 与输入进行残差链接，并进行层归一化，[batch_size, len_q, embedding_dim]
output = self.layer_norm(output + residual) # [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
# output [batch_size, len_q, embedding_dim]
#-------------------------维度信息--------------------------------
return output # 返回加入残差连接后层归一化的结果正弦位置编码表
Transformer使用Attention机制实现了网络的并行计算，这大大加快了网络的推理速度，但是也使得数据失去了时间上的先后性（RNN、LSTM本身就具有先后关系），所以需要引入一个位置编码表来融合位置信息。
# 生成正弦位置编码表的函数，用于在Transformer中引入位置信息
def get_sin_enc_table(n_position, embedding_dim):
#-------------------------维度信息--------------------------------
# n_position: 输入序列的最大长度
# embedding_dim: 词嵌入向量的维度
#-----------------------------------------------------------------
# 根据位置和维度信息，初始化正弦位置编码表
sinusoid_table = np.zeros((n_position, embedding_dim))
# 遍历所有位置和维度，计算角度值
for pos_i in range(n_position):
for hid_j in range(embedding_dim):
angle = pos_i / np.power(10000, 2 * (hid_j // 2) / embedding_dim)
sinusoid_table[pos_i, hid_j] = angle
# 计算正弦和余弦值
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 偶数维
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 奇数维
#-------------------------维度信息--------------------------------
# sinusoid_table 的维度是 [n_position, embedding_dim]
#-----------------------------------------------------------------
return torch.FloatTensor(sinusoid_table) # 返回正弦位置编码表填充注意力掩码
一个batch里面会有很多条句子，不同句子的长度可能会不一样，要让这些句子组成一个batch就需要在较短的句子后面填充PAD，使它们有相同的长度，但是这些PAD并不包含任何信息，所以在Transformer计算注意力的时候就需要给这些PAD打上马赛克（替换为一个非常小的数，在softmax步骤会被看成0）
# 生成填充注意力掩码的函数，用于在多头自注意力计算中忽略填充部分
def get_attn_pad_mask(seq_q, seq_k):
#-------------------------维度信息--------------------------------
# seq_q 的维度是 [batch_size, len_q]
# seq_k 的维度是 [batch_size, len_k]
#-----------------------------------------------------------------
batch_size, len_q = seq_q.size()
batch_size, len_k = seq_k.size()
# 生成布尔类型张量[batch_size，1，len_k(=len_q)]
pad_attn_mask = seq_k.data.eq(0).unsqueeze(1) #<PAD> Token的编码值为0
# 变形为何注意力分数相同形状的张量 [batch_size，len_q，len_k]
pad_attn_mask = pad_attn_mask.expand(batch_size, len_q, len_k)
#-------------------------维度信息--------------------------------
# pad_attn_mask 的维度是 [batch_size，len_q，len_k]
#-----------------------------------------------------------------
return pad_attn_mask # [batch_size，len_q，len_k]编码器层
一个多头自注意力+一个前馈神经网络，非常简单
# 定义编码器层类
class EncoderLayer(nn.Module):
def __init__(self):
super(EncoderLayer, self).__init__()
self.enc_self_attn = MultiHeadAttention() #多头自注意力层
self.pos_ffn = PoswiseFeedForwardNet() # 位置前馈神经网络层
def forward(self, enc_inputs, enc_self_attn_mask):
#-------------------------维度信息--------------------------------
# enc_inputs 的维度是 [batch_size, seq_len, embedding_dim]
# enc_self_attn_mask 的维度是 [batch_size, seq_len, seq_len]
#-----------------------------------------------------------------
# 将相同的Q，K，V输入多头自注意力层
enc_outputs, attn_weights = self.enc_self_attn(enc_inputs, enc_inputs,
enc_inputs, enc_self_attn_mask)
# 将多头自注意力outputs输入位置前馈神经网络层
enc_outputs = self.pos_ffn(enc_outputs)
#-------------------------维度信息--------------------------------
# enc_outputs 的维度是 [batch_size, seq_len, embedding_dim] 维度与 enc_inputs 相同
# attn_weights 的维度是 [batch_size, n_heads, seq_len, seq_len] 在注意力掩码维度上增加了头数
#-----------------------------------------------------------------
return enc_outputs, attn_weights # 返回编码器输出和每层编码器注意力权重编码器
编码器是6层编码器层的叠加，内部会完成位置编码的计算与融合、填充掩码的计算
# 定义编码器类
n_layers = 6 # 设置Encoder/Decoder的层数
class Encoder(nn.Module):
def __init__(self, corpus):
super(Encoder, self).__init__()
self.src_emb = nn.Embedding(corpus.src_vocab, d_embedding) # 词嵌入层
self.pos_emb = nn.Embedding.from_pretrained( \
get_sin_enc_table(corpus.src_len+1, d_embedding), freeze=True) # 位置嵌入层
self.layers = nn.ModuleList(EncoderLayer() for _ in range(n_layers))# 编码器层数
def forward(self, enc_inputs):
#-------------------------维度信息--------------------------------
# enc_inputs 的维度是 [batch_size, source_len]
#-----------------------------------------------------------------
# 创建一个从1到source_len的位置索引序列
pos_indices = torch.arange(1, enc_inputs.size(1) + 1).unsqueeze(0).to(enc_inputs)
#-------------------------维度信息--------------------------------
# pos_indices 的维度是 [1, source_len]
#-----------------------------------------------------------------
# 对输入进行词嵌入和位置嵌入相加 [batch_size, source_len，embedding_dim]
enc_outputs = self.src_emb(enc_inputs) + self.pos_emb(pos_indices)
#-------------------------维度信息--------------------------------
# enc_outputs 的维度是 [batch_size, seq_len, embedding_dim]
#-----------------------------------------------------------------
enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) # 生成自注意力掩码
#-------------------------维度信息--------------------------------
# enc_self_attn_mask 的维度是 [batch_size, len_q, len_k]
#-----------------------------------------------------------------
enc_self_attn_weights = [] # 初始化 enc_self_attn_weights
# 通过编码器层 [batch_size, seq_len, embedding_dim]
for layer in self.layers:
enc_outputs, enc_self_attn_weight = layer(enc_outputs, enc_self_attn_mask)
enc_self_attn_weights.append(enc_self_attn_weight)
#-------------------------维度信息--------------------------------
# enc_outputs 的维度是 [batch_size, seq_len, embedding_dim] 维度与 enc_inputs 相同
# enc_self_attn_weights 是一个列表，每个元素的维度是[batch_size, n_heads, seq_len, seq_len]
#-----------------------------------------------------------------
return enc_outputs, enc_self_attn_weights # 返回编码器输出和编码器注意力权重后续注意力掩码
在模型实际使用的时候，Transformer解码器和RNN/LSTM一样，接收到来自编码器的上下文向量后，从一个sos起始符开始输出文本，每输出一个词，这个词会作为输入重新输入到解码器中，输出下一个词，直到终止符eos出现，这个过程可以叫贪心迭代算法。
模型每次输出的句子都会比上一次输出的长一个单词，除了这个多出来的单词外别的地方都一样，如果我们把模型每次输出的句子放在一起，将相同内容对齐，可以得到一个直角三角形，这个三角形里的内容可以看成是训练时模型的已知信息，用来teacher forcing。
在训练模型的时候，Transformer为了进行并行计算并不会一次一次迭代，而是会将模型正确迭代时输入的内容作为teacher forcing喂进去，在解码器对输入进行自注意力计算时使用一个倒三角的注意力掩码，将模型未知的，或者说现阶段模型不应该知道的信息码住。
# 生成后续注意力掩码的函数，用于在多头自注意力计算中忽略未来信息
def get_attn_subsequent_mask(seq):
#-------------------------维度信息--------------------------------
# seq 的维度是 [batch_size, seq_len(Q)=seq_len(K)]
#-----------------------------------------------------------------
attn_shape = [seq.size(0), seq.size(1), seq.size(1)] # 获取输入序列的形状
#-------------------------维度信息--------------------------------
# attn_shape是一个一维张量 [batch_size, seq_len(Q), seq_len(K)]
#-----------------------------------------------------------------
# 使用numpy创建一个上三角矩阵（triu = triangle upper）
subsequent_mask = np.triu(np.ones(attn_shape), k=1)
#-------------------------维度信息--------------------------------
# subsequent_mask 的维度是 [batch_size, seq_len(Q), seq_len(K)]
#-----------------------------------------------------------------
# 将numpy数组转换为PyTorch张量，并将数据类型设置为byte（布尔值）
subsequent_mask = torch.from_numpy(subsequent_mask).byte()
#-------------------------维度信息--------------------------------
# 返回的subsequent_mask 的维度是 [batch_size, seq_len(Q), seq_len(K)]
#-----------------------------------------------------------------
return subsequent_mask # 返回后续位置的注意力掩码解码器层
与编码器层基本同理
# 定义解码器层类
class DecoderLayer(nn.Module):
def __init__(self):
super(DecoderLayer, self).__init__()
self.dec_self_attn = MultiHeadAttention() # 多头自注意力层
self.dec_enc_attn = MultiHeadAttention() # 多头注意力层，连接编码器和解码器
self.pos_ffn = PoswiseFeedForwardNet() # 位置前馈神经网络层
def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
#-------------------------维度信息--------------------------------
# dec_inputs 的维度是 [batch_size, target_len, embedding_dim]
# enc_outputs 的维度是 [batch_size, source_len, embedding_dim]
# dec_self_attn_mask 的维度是 [batch_size, target_len, target_len]
# dec_enc_attn_mask 的维度是 [batch_size, target_len, source_len]
#-----------------------------------------------------------------
# 将相同的Q，K，V输入多头自注意力层
dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs,
dec_inputs, dec_self_attn_mask)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
# dec_self_attn 的维度是 [batch_size, n_heads, target_len, target_len]
#-----------------------------------------------------------------
# 将解码器输出和编码器输出输入多头注意力层
dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs,
enc_outputs, dec_enc_attn_mask)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
# dec_enc_attn 的维度是 [batch_size, n_heads, target_len, source_len]
#-----------------------------------------------------------------
# 输入位置前馈神经网络层
dec_outputs = self.pos_ffn(dec_outputs)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
# dec_self_attn 的维度是 [batch_size, n_heads, target_len, target_len]
# dec_enc_attn 的维度是 [batch_size, n_heads, target_len, source_len]
#-----------------------------------------------------------------
# 返回解码器层输出，每层的自注意力和解-编编码器注意力权重
return dec_outputs, dec_self_attn, dec_enc_attn解码器
与编码器接近，要注意的是，编码器最终的输出会输入到全部的6层解码器中
# 定义解码器类
n_layers = 6 # 设置Decoder的层数
class Decoder(nn.Module):
def __init__(self, corpus):
super(Decoder, self).__init__()
self.tgt_emb = nn.Embedding(corpus.tgt_vocab, d_embedding) # 词嵌入层
self.pos_emb = nn.Embedding.from_pretrained( \
get_sin_enc_table(corpus.tgt_len+1, d_embedding), freeze=True) # 位置嵌入层
self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)]) # 叠加多层
def forward(self, dec_inputs, enc_inputs, enc_outputs):
#-------------------------维度信息--------------------------------
# dec_inputs 的维度是 [batch_size, target_len]
# enc_inputs 的维度是 [batch_size, source_len]
# enc_outputs 的维度是 [batch_size, source_len, embedding_dim]
#-----------------------------------------------------------------
# 创建一个从1到source_len的位置索引序列
pos_indices = torch.arange(1, dec_inputs.size(1) + 1).unsqueeze(0).to(dec_inputs)
#-------------------------维度信息--------------------------------
# pos_indices 的维度是 [1, target_len]
#-----------------------------------------------------------------
# 对输入进行词嵌入和位置嵌入相加
dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(pos_indices)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
#-----------------------------------------------------------------
# 生成解码器自注意力掩码和解码器-编码器注意力掩码
dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs) # 填充位掩码
dec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs) # 后续位掩码
dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask.to(device) \
+ dec_self_attn_subsequent_mask.to(device)), 0)
dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs) # 解码器-编码器掩码
#-------------------------维度信息--------------------------------
# dec_self_attn_pad_mask 的维度是 [batch_size, target_len, target_len]
# dec_self_attn_subsequent_mask 的维度是 [batch_size, target_len, target_len]
# dec_self_attn_mask 的维度是 [batch_size, target_len, target_len]
# dec_enc_attn_mask 的维度是 [batch_size, target_len, source_len]
#-----------------------------------------------------------------
dec_self_attns, dec_enc_attns = [], [] # 初始化 dec_self_attns, dec_enc_attns
# 通过解码器层 [batch_size, seq_len, embedding_dim]
for layer in self.layers:
dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs,
dec_self_attn_mask, dec_enc_attn_mask)
dec_self_attns.append(dec_self_attn)
dec_enc_attns.append(dec_enc_attn)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
# dec_self_attns 是一个列表，每个元素的维度是 [batch_size, n_heads, target_len, target_len]
# dec_enc_attns 是一个列表，每个元素的维度是 [batch_size, n_heads, target_len, source_len]
#-----------------------------------------------------------------
# 返回解码器输出，解码器自注意力和解-编编码器注意力权重
return dec_outputs, dec_self_attns, dec_enc_attns大的来了:Transformer
经过拆解后的Transformer非常简单
# 定义Transformer模型
class Transformer(nn.Module):
def __init__(self, corpus):
super(Transformer, self).__init__()
self.encoder = Encoder(corpus) # 初始化编码器实例
self.decoder = Decoder(corpus) # 初始化解码器实例
# 定义线性投影层，将解码器输出转换为目标词汇表大小的概率分布
self.projection = nn.Linear(d_embedding, corpus.tgt_vocab, bias=False)
def forward(self, enc_inputs, dec_inputs):
#-------------------------维度信息--------------------------------
# enc_inputs 的维度是 [batch_size, source_seq_len]
# dec_inputs 的维度是 [batch_size, target_seq_len]
#-----------------------------------------------------------------
# 将输入传递给编码器，并获取编码器输出和自注意力权重
enc_outputs, enc_self_attns = self.encoder(enc_inputs)
#-------------------------维度信息--------------------------------
# enc_outputs 的维度是 [batch_size, source_len, embedding_dim]
# enc_self_attns 是一个列表，每个元素的维度是 [batch_size, n_heads, src_seq_len, src_seq_len]
#-----------------------------------------------------------------
# 将编码器输出、解码器输入和编码器输入传递给解码器
# 获取解码器输出、解码器自注意力权重和编码器-解码器注意力权重
dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)
#-------------------------维度信息--------------------------------
# dec_outputs 的维度是 [batch_size, target_len, embedding_dim]
# dec_self_attns 是一个列表，每个元素的维度是 [batch_size, n_heads, tgt_seq_len, src_seq_len]
# dec_enc_attns 是一个列表，每个元素的维度是 [batch_size, n_heads, tgt_seq_len, src_seq_len]
#-----------------------------------------------------------------
# 将解码器输出传递给投影层，生成目标词汇表大小的概率分布
dec_logits = self.projection(dec_outputs)
#-------------------------维度信息--------------------------------
# dec_logits 的维度是 [batch_size, tgt_seq_len, tgt_vocab_size]
#-----------------------------------------------------------------
# 返回逻辑值(原始预测结果),编码器自注意力权重，解码器自注意力权重，解-编码器注意力权重
return dec_logits, enc_self_attns, dec_self_attns, dec_enc_attnsTransformer机器翻译
深蓝学院 生成式预训练语言模型 第七章作业
训练
训练数据集构建
数据集用的是是深蓝学院给的课程资料，这里就不放出来了;导入的CorpusLoader是黄佳老师提供的一个工具，可以通过它快速地构建一个词表，需要注意的是：每次运行构建的词表是不一样的，所以我们需要把保存两个dic，部署预训练模型的时候导入
from CorpusLoader import CorpusLoader
from torch.utils.data import DataLoader
import torch.nn as nn
corpus_loader = CorpusLoader('all_sentences.txt')
corpus_loader.process_sentences()
corpus_loader.build_vocab()
# 保存词表
import pickle
with open("corpus_loader.word2idx_cn.pkl", "wb") as tf:
pickle.dump(corpus_loader.word2idx_cn,tf)
with open("corpus_loader.idx2word_en.pkl", "wb") as tf:
pickle.dump(corpus_loader.idx2word_en,tf)
def collate_fn(batch):
batch.sort(key=lambda x: len(x[0]), reverse=True)
sentence_cn, sentence_en_in, sentence_en_out = zip(*batch)
sentence_cn = nn.utils.rnn.pad_sequence(sentence_cn, padding_value=corpus_loader.word2idx_cn['<pad>'],batch_first=True)
sentence_en_in = nn.utils.rnn.pad_sequence(sentence_en_in, padding_value=corpus_loader.word2idx_en['<pad>'],batch_first=True)
sentence_en_out = nn.utils.rnn.pad_sequence(sentence_en_out, padding_value=corpus_loader.word2idx_en['<pad>'],batch_first=True)
return sentence_cn, sentence_en_in, sentence_en_out
dataset = corpus_loader.create_dataset()
dataloader = DataLoader(dataset, batch_size=16, shuffle=True, collate_fn=collate_fn)训练循环
Transformer还真挺大的，搁我自己的电脑上跑太慢了，我去网上租了一个3080服务器跑了50个循环，用了差不多4个小时
import torch # 导入torch
import torch.optim as optim # 导入优化器
from Transformer_Model import Transformer
device = "cuda" if torch.cuda.is_available() else "cpu"
# 创建模型实例并选择优化器
model = Transformer(corpus_loader).to(device)
model.train()
criterion = torch.nn.CrossEntropyLoss(ignore_index=corpus_loader.word2idx_en['<pad>']) # 忽略padding的损失
optimizer = optim.Adam(model.parameters(), lr=0.0001)
# 进行训练
loss_cache = 9999
for epoch in range(50):
for i, (enc_inputs, dec_inputs, target_batch) in enumerate(dataloader):
enc_inputs, dec_inputs, target_batch = enc_inputs.to(device), dec_inputs.to(device), target_batch.to(device)
optimizer.zero_grad()
outputs, _, _, _ = model(enc_inputs, dec_inputs)
loss = criterion(outputs.view(-1, outputs.size(-1)), target_batch.view(-1))
loss.backward()
optimizer.step()
if (i+1) % 100 == 0:
print(f"Epoch[{epoch+1}/50], Step[{i+1}/{len(dataloader)}], Loss: {loss.item()}")
if (epoch + 1) % 1 == 0: # 打印损失并保存模型参数
print(f"Epoch: {epoch + 1:04d} cost = {loss:.6f}")
if float(loss.item()) < loss_cache:
print("保存模型")
loss_cache = loss
torch.save(model.state_dict(), 'transformer.pth')模型部署
import torch # 导入torch
from CorpusLoader import CorpusLoader
corpus_loader = CorpusLoader('all_sentences.txt')
corpus_loader.process_sentences()
corpus_loader.build_vocab()
# 导入词表
import pickle
with open("corpus_loader.idx2word_en.pkl", "rb") as tf:
idx2word_en = pickle.load(tf)
with open("corpus_loader.word2idx_cn.pkl", "rb") as tf:
word2idx_cn = pickle.load(tf)
from Transformer_Model import Transformer
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Transformer(corpus_loader).to(device)
model.load_state_dict(torch.load('transformer.pth'))
model.eval()
# 贪心算法迭代函数
def greedy_decoder(model, enc_inputs, start_symbol):
# 将输入传递给编码器，并获取编码器输出和自注意力权重
enc_outputs, enc_self_attns = model.encoder(enc_inputs)
dec_input = torch.zeros(1,1).type_as(enc_inputs.data)
dec_input[0][0] = start_symbol
for i in range(100):# 限制模型的最长输出为99词
# 将解码器输入、编码器输入输出传递给解码器，获取解码器输出、自注意力权重和编-解码器注意力权重
dec_outputs, _, _ = model.decoder(dec_input, enc_inputs, enc_outputs)
# 将解码器输出传递给投影层，生成目标词汇表大小的概率分布
predict = model.projection(dec_outputs)
# 找到概率最大的下一个词
prob = predict.squeeze(0).max(dim=-1, keepdim=False)[1][-1]
# 拼接到dec_input
dec_input_cache = torch.zeros(1,1).type_as(enc_inputs.data)
dec_input_cache[0][0] = prob
dec_input = torch.cat((dec_input,dec_input_cache), dim=1)
if int(prob) == 0:# 如果输出为eos，中断循环
break
return dec_input
import jieba
while True:
sentence_cn = input("翻译器输入: ")
sentence_cn = ' '.join(jieba.cut(sentence_cn.strip(), cut_all=False)).split(" ")# 使用jieba分词
sentence_idx = [word2idx_cn[i] for i in sentence_cn]
sentence_idx = torch.tensor(sentence_idx).to(device).unsqueeze(0)
enc_inputs = sentence_idx
greedy_dec_output = greedy_decoder(model, enc_inputs, start_symbol=1)# 获得贪婪解码器输入
greedy_dec_output_words = [idx2word_en[int(i)] for i in greedy_dec_output.squeeze()]# 转换为单词序列
print(" ".join(greedy_dec_output_words[1:-1]))
print("")这时候模型就可以跑了，不知道为什么翻译的效果不是特别好，如果输入的中文句子稍微长一点模型还算能说出人话，但是如果稍微短一些（如下图所示），这个Transformer自己在没话找话
BLEU得分
重新修改训练部分的代码，这一次在每训练一轮后将BLEU得分用tensorboard记录，考虑到之前训练50轮的效果不是很好，这次训练时间长一点，同时把loss也用tensorboard记录一下
# 计算BLEU得分
from torchtext.data.metrics import bleu_score
def compute_bleu(model, dataset):
model.eval()
total_score = 0.0
with torch.no_grad():
for i in range(dataset.__len__()):
sentence_cn, _, sentence_en_out = dataset.__getitem__(i)
sentence_cn = torch.tensor(sentence_cn).to(device).unsqueeze(0)
sentence_en_out = [corpus_loader.idx2word_en[int(i)] for i in sentence_en_out]
#迭代推理
dec_input = greedy_decoder(model, sentence_cn, start_symbol=1)
greedy_dec_output_words = [corpus_loader.idx2word_en[int(i)] for i in dec_input.squeeze()]# 转换为单词序列
candidate_corpus = [greedy_dec_output_words[1:]]
references_corpus = [[sentence_en_out]]
# 计算BLEU得分
total_score += bleu_score(candidate_corpus, references_corpus)
#print(i)
if i > 500:# 全部算一遍太慢了，这里就只针对前500个左右算吧
break
return total_score / 502在我斥巨资算了14个半小时后，模型得分和loss如下：
参考
深蓝学院黄佳老师代码
论文Attention Is All You Need