Transformer 模型是采用自注意力机制的深度学习模型，允许更多的并行计算，以此减少训练时间。下述代码是在参考网上分享的的代码和文章后，使用 PyTorch 实现的 Transformer 模型。
完整代码
import torch
import torch.nn as nn
import math
import numpy as np
from typing import List, Dict, Tuple
PAD_WORD = "<PAD>"
UNK_WORD = "<UNK>"
BOS_WORD = "<SOS>"
EOS_WORD = "<EOS>"
class PositionalEncoding(nn.Module):
r"""
为输入序列中的每个位置添加独特的编码向量。
"""
def __init__(self, embedding_size: int, max_sequence_length: int):
"""
Args:
embedding_size: 词向量的维度
max_sequence_length: 最长序列长度
"""
super(PositionalEncoding, self).__init__()
# 生成正弦位置编码表
sinusoid_table = PositionalEncoding.__get_sinusoid_encoding_table(embedding_size, max_sequence_length)
self.register_buffer("sinusoid_table", sinusoid_table)
def forward(self, x: torch.tensor) -> torch.Tensor:
# 获取张量 x 对应长度的位置编码
sinusoid_encoding = self.sinusoid_table[:, :x.size(1)]
# 对张量 x 与位置编码做相加操作
output = x + sinusoid_encoding
return output
@staticmethod
def __get_sinusoid_encoding_table(embedding_size: int, max_sequence_length: int) -> torch.Tensor:
sinusoid_list = [PositionalEncoding.__get_position_angle_vec(embedding_size, curr_position)
for curr_position in range(max_sequence_length)]
sinusoid_ndarray = np.array(sinusoid_list)
sinusoid_ndarray[:, 0::2] = np.sin(sinusoid_ndarray[:, 0::2])
sinusoid_ndarray[:, 1::2] = np.cos(sinusoid_ndarray[:, 1::2])
sinusoid_tensor = torch.FloatTensor(sinusoid_ndarray)
sinusoid_tensor = sinusoid_tensor.unsqueeze(0)
return sinusoid_tensor
@staticmethod
def __get_position_angle_vec(embedding_size: int, curr_position: int) -> List[float]:
result = [curr_position / math.pow(10000, 2 * (curr_dim // 2) / embedding_size)
for curr_dim in range(embedding_size)]
return result
class PositionWiseFeedForward(nn.Module):
r"""
编码器和解码器的每层都包含一个全连接前馈网络。
"""
def __init__(self, embedding_size: int, hidden_size: int, dropout_prob=0.1):
"""
Args:
embedding_size: 词向量的维度
hidden_size: 隐藏层维度
dropout_prob: 丢弃率, 防止过拟合
"""
super(PositionWiseFeedForward, self).__init__()
self.linear_1 = nn.Linear(embedding_size, hidden_size)
self.linear_2 = nn.Linear(hidden_size, embedding_size)
self.layer_norm = nn.LayerNorm(embedding_size, eps=1e-6)
self.dropout = nn.Dropout(dropout_prob)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
output = self.linear_1(x)
output = nn.functional.relu(output)
# 通过丢弃法防止过拟合
output = self.dropout(output)
output = self.linear_2(output)
# 做残差连接操作
output = output + residual
# 做归一化操作
output = self.layer_norm(output)
return output
class ScaledDotProductAttention(nn.Module):
r"""
缩放点积注意力通过计算查询 Q 与键 K 之间的相似度来分配注意力权重。
"""
def __init__(self, key_dim: int, dropout_prob=0.1):
"""
Args:
key_dim: key维度
dropout_prob: 丢弃率, 防止过拟合
"""
super(ScaledDotProductAttention, self).__init__()
self.key_dim = key_dim
self.sqrt_of_key_dim = math.sqrt(key_dim)
self.dropout = nn.Dropout(dropout_prob)
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
attention_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# 转置前 key 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, key维度),
# 转置后 key_t 的 shape 属性是 (批次大小, 多头注意力的数量, key维度, 序列长度)
key_t = torch.transpose(key, 2, 3)
# 张量 query 与 key_t 的矩阵乘法操作对应注意力算法中 query 与 key 做向量点积操作，
# 其结果 attention_score 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, 序列长度)
attention_score = torch.matmul(query, key_t) / self.sqrt_of_key_dim
# 通过掩码遮挡用于统一长度的填充标识符以及在解码阶段遮挡未预测信息，
# 其中编码器自注意力掩码的 shape 属性是 (批次大小, 1, 1, 序列长度),
# 其中解码器自注意力掩码的 shape 属性是 (批次大小, 1, 序列长度, 序列长度)
if attention_mask is not None:
attention_score = torch.masked_fill(attention_score, mask=attention_mask,
value=-1e9)
attention_score = nn.functional.softmax(attention_score, dim=-1)
# 通过丢弃法防止过拟合
attention_score = self.dropout(attention_score)
# 注意力分数与张量 value 的矩阵乘法操作对应注意力算法中注意力分数与 value 做向量点积操作，
# 其结果 output 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, value维度)
output = torch.matmul(attention_score, value)
return output, attention_score
class MultiHeadAttention(nn.Module):
r"""
多头注意力通过并行地运行多个独立的注意力机制来获取输入序列的不同子空间的注意力分布，
从而更全面地捕获序列中潜在的多种语义关联。
"""
def __init__(self, head_num: int, embedding_size: int, dropout_prob=0.1):
"""
Args:
head_num: 多头注意力的数量
embedding_size: 词向量的维度
dropout_prob: 丢弃率, 防止过拟合
"""
super(MultiHeadAttention, self).__init__()
self.head_num = head_num
self.embedding_size = embedding_size
self.query_linear = nn.Linear(embedding_size, embedding_size, bias=False)
self.key_linear = nn.Linear(embedding_size, embedding_size, bias=False)
self.value_linear = nn.Linear(embedding_size, embedding_size, bias=False)
self.out_linear = nn.Linear(embedding_size, embedding_size, bias=False)
self.layer_norm = nn.LayerNorm(embedding_size, eps=1e-6)
self.query_dim = embedding_size // head_num
self.key_dim = embedding_size // head_num
self.value_dim = embedding_size // head_num
# 缩放点积注意力
self.scaled_dot_product_attention = ScaledDotProductAttention(key_dim=self.key_dim,
dropout_prob=dropout_prob)
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
attention_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
residual = query
# 批次大小
batch_size = query.size(0)
# 序列长度
sequence_length = query.size(1)
# 操作后 query 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
query = self.query_linear(query)
key = self.key_linear(key)
value = self.value_linear(value)
# 维度修改后 query 的 shape 属性是 (批次大小, 序列长度, 多头注意力的数量, query维度)
query = query.view(batch_size, -1, self.head_num, self.query_dim)
# 转置后 query 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, query维度)
query = torch.transpose(query, 1, 2)
key = key.view(batch_size, -1, self.head_num, self.key_dim)
key = torch.transpose(key, 1, 2)
value = value.view(batch_size, -1, self.head_num, self.value_dim)
value = torch.transpose(value, 1, 2)
if attention_mask is not None:
# 其中编码器自注意力掩码解压前 attention_mask 的 shape 属性是 (批次大小, 1, 序列长度)，
# 解压后 attention_mask 的 shape 属性是 (批次大小, 1, 1, 序列长度)。
# 其中解码器自注意力掩码解压前 attention_mask 的 shape 属性是 (批次大小, 序列长度, 序列长度)，
# 解压后 attention_mask 的 shape 属性是 (批次大小, 1, 序列长度, 序列长度)。
attention_mask = torch.unsqueeze(attention_mask, 1)
# 缩放点积注意力返回结果 output 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, value维度)
output, attention_score = self.scaled_dot_product_attention(query, key, value, attention_mask)
# 转置后 output 的 shape 属性是 (批次大小, 序列长度, 多头注意力的数量, value维度)
output = torch.transpose(output, 1, 2)
output = output.contiguous()
# 维度修改后 output 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
output = output.view(batch_size, sequence_length, -1)
output = self.out_linear(output)
# 做残差连接操作
output = output + residual
# 做归一化操作
output = self.layer_norm(output)
return output, attention_score
class EncoderLayer(nn.Module):
r"""
通过堆叠多个相同的编码器层组成编码器，其中每层包括两个子层分别是 多头自注意力、前馈神经网络。
"""
def __init__(self, self_attention: MultiHeadAttention, feed_forward: PositionWiseFeedForward):
"""
Args:
decoder_self_attention: 自注意力
feed_forward: 前馈神经网络
"""
super(EncoderLayer, self).__init__()
self.self_attention = self_attention
self.feed_forward = feed_forward
def forward(self, encoder_inputs: torch.Tensor, self_attention_mask: torch.Tensor) \
-> Tuple[torch.Tensor, torch.Tensor]:
# 自注意力的返回结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
encoder_outputs, attention_score = self.self_attention(query=encoder_inputs,
key=encoder_inputs,
value=encoder_inputs,
attention_mask=self_attention_mask)
# 前馈神经网络的返回结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
encoder_outputs = self.feed_forward(encoder_outputs)
return encoder_outputs, attention_score
class DecoderLayer(nn.Module):
r"""
通过堆叠多个相同的解码器层组成编码器，其中每层包括三个子层分别是 多头自注意力、解码器-编码器注意力、前馈神经网络。
"""
def __init__(self, decoder_self_attention: MultiHeadAttention, decoder_encoder_attention: MultiHeadAttention,
feed_forward: PositionWiseFeedForward):
"""
Args:
decoder_self_attention: 自注意力
decoder_encoder_attention: 解码器-编码器注意力
feed_forward: 前馈神经网络
"""
super(DecoderLayer, self).__init__()
self.decoder_self_attention = decoder_self_attention
self.decoder_encoder_attention = decoder_encoder_attention
self.feed_forward = feed_forward
def forward(self, decoder_inputs: torch.Tensor, encoder_outputs: torch.Tensor,
dec_self_attn_mask: torch.Tensor, dec_enc_attn_mask: torch.Tensor) \
-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# 自注意力的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
decoder_outputs, dec_self_attn_score = self.decoder_self_attention(query=decoder_inputs,
key=decoder_inputs,
value=decoder_inputs,
attention_mask=dec_self_attn_mask)
# 解码器-编码器注意力的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
decoder_outputs, dec_enc_attn_score = self.decoder_encoder_attention(query=decoder_outputs,
key=encoder_outputs,
value=encoder_outputs,
attention_mask=dec_enc_attn_mask)
# 前馈神经网络的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
decoder_outputs = self.feed_forward(decoder_outputs)
return decoder_outputs, dec_self_attn_score, dec_enc_attn_score
class Encoder(nn.Module):
r"""
编码器的任务是将输入序列转换成高维向量表示。
"""
def __init__(self, encoder_layers: nn.ModuleList, word_embedding: nn.Embedding,
positional_encoding: PositionalEncoding, dropout_prob=0.1):
"""
Args:
encoder_layers: 编码器层
word_embedding: 词向量表
positional_encoding: 位置编码
dropout_prob: 丢弃率, 防止过拟合
"""
super(Encoder, self).__init__()
self.encoder_layers = encoder_layers
self.word_embedding = word_embedding
self.positional_encoding = positional_encoding
self.dropout = nn.Dropout(dropout_prob)
def forward(self, encoder_inputs: torch.Tensor, self_attention_mask: torch.Tensor) \
-> Tuple[torch.Tensor, torch.Tensor]:
self_attn_score_list = []
# 词嵌入后变量 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
encoder_outputs = self.word_embedding(encoder_inputs)
# 添加位置编码
encoder_outputs = self.positional_encoding(encoder_outputs)
# 通过丢弃法防止过拟合
encoder_outputs = self.dropout(encoder_outputs)
# 依次把输出传递到下一个编码器层，其中 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
for encoder_layer in self.encoder_layers:
encoder_outputs, self_attn_score = encoder_layer(encoder_outputs, self_attention_mask)
self_attn_score_list.append(self_attn_score)
return encoder_outputs, self_attn_score_list
class Decoder(nn.Module):
r"""
解码器的任务是根据编码器生成的高维向量表示和已经生成的目标序列，逐步生成新的序列。
"""
def __init__(self, decoder_layers: nn.ModuleList, word_embedding: nn.Embedding,
positional_encoding: PositionalEncoding, dropout_prob=0.1):
"""
Args:
decoder_layers: 编码器层
word_embedding: 词向量表
positional_encoding: 位置编码
dropout_prob: 丢弃率, 防止过拟合
"""
super(Decoder, self).__init__()
self.decoder_layers = decoder_layers
self.word_embedding = word_embedding
self.positional_encoding = positional_encoding
self.dropout = nn.Dropout(dropout_prob)
def forward(self, decoder_inputs: torch.Tensor, encoder_outputs: torch.Tensor,
dec_self_attn_mask: torch.Tensor, dec_enc_attn_mask: torch.Tensor) \
-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
dec_self_attn_score_list = []
dec_enc_attn_score_list = []
# 词嵌入后变量 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
decoder_outputs = self.word_embedding(decoder_inputs)
# 添加位置编码
decoder_outputs = self.positional_encoding(decoder_outputs)
# 通过丢弃法防止过拟合
decoder_outputs = self.dropout(decoder_outputs)
# 依次把输出传递到下一个解码器层，其中 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
for decoder_layer in self.decoder_layers:
decoder_outputs, dec_self_attn_score, dec_enc_attn_score = decoder_layer(decoder_inputs=decoder_outputs,
encoder_outputs=encoder_outputs,
dec_self_attn_mask=dec_self_attn_mask,
dec_enc_attn_mask=dec_enc_attn_mask)
dec_self_attn_score_list.append(dec_self_attn_score)
dec_enc_attn_score_list.append(dec_enc_attn_score)
return decoder_outputs, dec_self_attn_score_list, dec_enc_attn_score_list
class Transformer(nn.Module):
r"""
Transformer模型是基于自注意力机制的深度学习模型。
"""
def __init__(self, encoder: Encoder, decoder: Decoder, output_projection: nn.Module, pad_index: int):
"""
Args:
encoder: 编码器
decoder: 解码器
output_projection: 输出向量投影到词表空间
pad_index: 填充标识符索引
"""
super(Transformer, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.output_projection = output_projection
self.pad_index = pad_index
def forward(self, encoder_inputs: torch.Tensor, decoder_inputs: torch.Tensor,
encoder_outputs=None) -> torch.Tensor:
# 变量 encooder_mask 的 shape 属性是 (批次大小, 1, 序列长度)
encooder_mask = self.__get_pad_mask(encoder_inputs, self.pad_index)
# 变量 decoder_mask 的 shape 属性是 (批次大小, 序列长度, 序列长度)
decoder_mask = self.__get_pad_mask(decoder_inputs, self.pad_index) | self.__get_subsequent_mask(decoder_inputs)
# 用于在推理阶段避免编码器做重复的计算
if encoder_outputs is None:
# 编码器输出结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
encoder_outputs, *_ = self.encoder(encoder_inputs, encooder_mask)
# 解码器输出结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
encoder_outputs, *_ = self.decoder(decoder_inputs, encoder_outputs, decoder_mask, encooder_mask)
# 投影到词表空间后 decoder_logit 的 shape 属性是 (批次大小, 序列长度, 解码器的词汇表大小)
decoder_logit = self.output_projection(encoder_outputs)
return decoder_logit, encoder_outputs
@staticmethod
def __get_pad_mask(x: torch.Tensor, pad_index: int) -> torch.Tensor:
# 变量 mask_tensor 的 shape 属性是 (批次大小, 序列长度)
mask_tensor = (x == pad_index)
# 解压后 mask_tensor 的 shape 属性是 (批次大小, 序列长度, 序列长度)
mask_tensor = mask_tensor.unsqueeze(-2)
return mask_tensor
@staticmethod
def __get_subsequent_mask(x: torch.Tensor) -> torch.Tensor:
_, sequence_length = x.size()
ones_tensor = torch.ones(1, sequence_length, sequence_length, device=x.device)
# 矩阵主对角线及之下区域的元素置为0
triu_tensor = torch.triu(ones_tensor, diagonal=1)
triu_tensor = triu_tensor.type(torch.uint8)
# 变量 mask_tensor 的 shape 属性是 (批次大小, 序列长度, 序列长度)
mask_tensor = (triu_tensor == 1)
return mask_tensor
测试代码
# 多头注意力的数量
HEAD_NUM = 8
# 编码器/解码器层的数量
LAYER_NUM = 6
# 序列最大长度
MAX_SEQUENCE_SIZE = 128
# 词汇表最大数目
MAX_VOCAB_SIZE = 10000
# 词向量的维度
EMBEDDING_SIZE = 512
# 前馈全连接层的隐藏层维度
HIDDEN_SIZE = 2048
def create_model(max_sequence_length):
positional_encoding = PositionalEncoding(EMBEDDING_SIZE, MAX_SEQUENCE_SIZE)
encoder_embedding = nn.Embedding(MAX_VOCAB_SIZE, EMBEDDING_SIZE, 0)
decoder_embedding = nn.Embedding(MAX_VOCAB_SIZE, EMBEDDING_SIZE, 0)
output_projection = nn.Linear(EMBEDDING_SIZE, max_sequence_length, bias=False)
encoder_layers = nn.ModuleList()
for i in range(LAYER_NUM):
self_attention = MultiHeadAttention(HEAD_NUM, EMBEDDING_SIZE)
feed_forward = PositionWiseFeedForward(EMBEDDING_SIZE, HIDDEN_SIZE)
encoder_layer = EncoderLayer(self_attention, feed_forward)
encoder_layers.append(encoder_layer)
decoder_layers = nn.ModuleList()
for i in range(LAYER_NUM):
decoder_self_attention = MultiHeadAttention(HEAD_NUM, EMBEDDING_SIZE)
decoder_encoder_attention = MultiHeadAttention(HEAD_NUM, EMBEDDING_SIZE)
feed_forward = PositionWiseFeedForward(EMBEDDING_SIZE, HIDDEN_SIZE)
decoder_layer = DecoderLayer(decoder_self_attention, decoder_encoder_attention, feed_forward)
decoder_layers.append(decoder_layer)
encoder = Encoder(encoder_layers, encoder_embedding, positional_encoding)
decoder = Decoder(decoder_layers, decoder_embedding, positional_encoding)
model = Transformer(encoder, decoder, output_projection, 0)
return model
def train_model(model, encoder_inputs, decoder_inputs, target_labels, device):
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
for epoch in range(20):
encoder_inputs = encoder_inputs.to(device)
decoder_inputs = decoder_inputs.to(device)
target_labels = target_labels.to(device)
optimizer.zero_grad()
outputs, *_ = model(encoder_inputs, decoder_inputs)
outputs = outputs.transpose(1, 2)
# 计算训练的误差
loss = criterion(outputs, target_labels)
print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))
# 根据误差反向计算梯度和更新权重
loss.backward()
optimizer.step()
return model
source_word2id_dict = {'<PAD>': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4}
target_word2id_dict = {'<PAD>': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, '<SOS>': 5, '<EOS>': 6}
target_id2word_dict = {id: word for word, id in target_word2id_dict.items()}
encoder_inputs = torch.LongTensor([[source_word2id_dict[word] for word in 'ich mochte ein bier <PAD>'.split()]])
decoder_inputs = torch.LongTensor([[target_word2id_dict[word] for word in '<SOS> i want a beer'.split()]])
target_labels = torch.LongTensor([[target_word2id_dict[word] for word in 'i want a beer <EOS>'.split()]])
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
# 创建模型
model = create_model(max_sequence_length=len(target_word2id_dict))
# 训练模型
train_model(model, encoder_inputs, decoder_inputs, target_labels, device=device)
# 测试模型
predict, *_ = model(encoder_inputs.to(device), decoder_inputs.to(device))
predict = predict.data.max(2, keepdim=False)[1]
print('ich mochte ein bier <PAD>', '->', [target_id2word_dict[id] for id in predict[0].tolist()])
