前置知识：嵌入层和位置编码、编码器的实现、解码器和输出部分的实现。
之前的文章已经把构建 Transformer 所需的所有组件构建完了，这篇文章开始构建整个编码器-解码器结构。
1 编码器-解码器的代码实现
python
class EncoderDecoder(nn.Module):
def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
"""
encoder: 编码器模块，用于处理源序列
decoder: 解码器模块，用于生成目标序列
src_embed: 源序列嵌入层，将源序列tokens转换为向量表示
tgt_embed: 目标序列嵌入层，将目标序列tokens转换为向量表示
generator: 生成器模块，将解码器输出转换为目标词汇表上的概率分布
"""
super(EncoderDecoder, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
self.tgt_embed = tgt_embed
self.generator = generator
def forward(self, src, tgt, src_mask, tgt_mask):
"""
Args:
src: 源序列，形状为 [batch_size, src_seq_len]
tgt: 目标序列，形状为 [batch_size, tgt_seq_len]
src_mask: 源序列的掩码，形状为 [batch_size, 1, src_seq_len]，用于屏蔽填充位置
tgt_mask: 目标序列的掩码，形状为 [batch_size, tgt_seq_len, tgt_seq_len]，用于屏蔽填充位置和未来位置
Return:
解码器的输出，形状为 [batch_size, tgt_seq_len, d_model]
"""
# 先对源序列进行编码，得到编码器输出
# 然后将编码器输出、源序列掩码、目标序列和目标序列掩码传递给解码器
return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
def encode(self, src, src_mask):
"""
Args:
src: 源序列，形状为 [batch_size, src_seq_len]
src_mask: 源序列的掩码，形状为 [batch_size, 1, src_seq_len]，用于屏蔽填充位置
Return:
编码器的输出，形状为 [batch_size, src_seq_len, d_model]，作为解码器的输入
"""
# (1) 对源序列进行嵌入处理，将tokens转换为向量表示
# (2) 将嵌入结果和源序列掩码传递给编码器
# (3) 返回编码器的输出，作为解码器的"记忆"输入
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
"""
Args:
memory: 编码器的输出
src_mask: 源序列的掩码
tgt: 目标序列
tgt_mask: 目标序列的掩码
"""
# (1) 对目标序列进行嵌入处理，将tokens转换为向量表示
# (2) 将嵌入结果、编码器输出、源序列掩码、目标序列掩码传给解码器
# (3) 返回解码器的输出
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)实例化验证一下代码实现
python
vocab_size = 1000
d_model = 512
encoder = en
decoder = de
source_embed = nn.Embedding(vocab_size, d_model)
target_embed = nn.Embedding(vocab_size, d_model)
generator = gen
# 假设源数据与目标数据相同，实际中并不相同
source = target = Variable(torch.LongTensor([[100, 2, 421, 508], [491, 998, 1, 221]]))
# 假设src_mask与tgt_mask相同，实际中并不相同
source_mask = target_mask = Variable(torch.zeros(2, 4, 4))
ed = EncoderDecoder(encoder, decoder, source_embed, target_embed, generator)
ed_result = ed(source, target, source_mask, target_mask)
print(ed_result)
print(ed_result.shape)运行结果
plaintext
tensor([[[-0.6909, -0.9724, 0.6361, ..., 0.0525, 0.5336, -0.4723],
[-1.7389, -1.2510, 0.7788, ..., -0.0156, 0.0258, -0.9280],
[-1.7369, -1.2625, 0.7260, ..., -0.4470, -1.0672, 0.1064],
[-1.3041, -1.3951, -0.0759, ..., -0.1239, -0.6066, -0.8090]],
[[-0.0288, 0.2606, 1.1938, ..., 0.3468, -0.5532, -1.2409],
[-0.2443, -0.1397, 0.7443, ..., -0.2610, 0.6653, -1.7663],
[-0.9244, -0.9061, 1.7313, ..., 1.1313, 1.1599, -2.4360],
[ 0.0872, -0.6195, 0.7045, ..., -0.0286, 0.4245, -0.6196]]],
grad_fn=<AddBackward0>)
torch.Size([2, 4, 512])2 完整模型构建
基于上述编码器-解码器结构构建用于训练的、完整的 Transformer 模型。
python
def make_model(
src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):
"""
Args:
src_vocab: 源语言词汇表大小
tgt_vocab: 目标语言词汇表大小
N: 编码器和解码器的层数
d_model: 模型的隐藏维度
d_ff: 前馈网络的隐藏层维度
h: 多头注意力的头数
dropout: Dropout 丢弃率
"""
# 创建一个深拷贝函数，用于复制组件实例，避免共享参数
c = copy.deepcopy
# 初始化多头注意力机制，处理不同子空间的注意力计算
attn = MultiHeadedAttention(h, d_model)
# 初始化位置前馈网络，对每个位置的表示进行非线性变换
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
# 初始化位置编码，为输入序列添加位置信息
position = PositionalEncoding(d_model, dropout)
# 构建完整模型
model = EncoderDecoder(
# 编码器由 N 个 EncoderLayer 组成，每个层包含自注意力机制(attn)和前馈网络(ff)
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
# 解码器由 N 个 DecoderLayer 组成，每个层包含自注意力机制(attn)、编码器-解码器注意力机制(attn)和前馈网络(ff)
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
# 源嵌入：将源语言词汇转换为词向量，并添加位置编码
nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
# 目标嵌入：将目标语言词汇转换为词向量，并添加位置编码
nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
# 生成器：将解码器输出转换为目标词汇表的概率分布
Generator(d_model, tgt_vocab),
)
# 参数初始化
# 使用 Xavier 均匀分布初始化所有维度大于 1 的参数，
# 这有助于模型的稳定训练和收敛
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
return model验证一下完整模型的代码实现
python
source_vocab = 11
target_vocab = 11
N = 6
res = make_model(source_vocab, target_vocab, N)
print(res)运行结果
plaintext
EncoderDecoder(
(encoder): Encoder(
(layers): ModuleList(
(0-5): 6 x EncoderLayer(
(self_attn): MultiHeadedAttention(
(linears): ModuleList(
(0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
)
(dropout): Dropout(p=0.1, inplace=False)
)
(feed_forward): PositionwiseFeedForward(
(w_1): Linear(in_features=512, out_features=2048, bias=True)
(w_2): Linear(in_features=2048, out_features=512, bias=True)
(dropout): Dropout(p=0.1, inplace=False)
)
(sublayer): ModuleList(
(0-1): 2 x SublayerConnection(
(norm): LayerNorm()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
)
(norm): LayerNorm()
)
(decoder): Decoder(
(layers): ModuleList(
(0-5): 6 x DecoderLayer(
(self_attn): MultiHeadedAttention(
(linears): ModuleList(
(0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
)
(dropout): Dropout(p=0.1, inplace=False)
)
(src_attn): MultiHeadedAttention(
(linears): ModuleList(
(0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
)
(dropout): Dropout(p=0.1, inplace=False)
)
(feed_forward): PositionwiseFeedForward(
(w_1): Linear(in_features=512, out_features=2048, bias=True)
(w_2): Linear(in_features=2048, out_features=512, bias=True)
(dropout): Dropout(p=0.1, inplace=False)
)
(sublayer): ModuleList(
(0-2): 3 x SublayerConnection(
(norm): LayerNorm()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
)
(norm): LayerNorm()
)
(src_embed): Sequential(
(0): Embeddings(
(lut): Embedding(11, 512)
)
(1): PositionalEncoding(
(dropout): Dropout(p=0.1, inplace=False)
)
)
(tgt_embed): Sequential(
(0): Embeddings(
(lut): Embedding(11, 512)
)
(1): PositionalEncoding(
(dropout): Dropout(p=0.1, inplace=False)
)
)
(generator): Generator(
(proj): Linear(in_features=512, out_features=11, bias=True)
)
)3 推理测试的代码实现
python
def inference_test():
# 创建一个小型 Transformer 模型，词汇表大小为 11，编码器和解码器各 2 层
test_model = make_model(11, 11, 2)
# 将模型设置为评估模式
test_model.eval()
# 创建一个长度为 10 的源序列
src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
# 创建一个全 1 的源掩码，表示所有位置都是有效的（非填充）
src_mask = torch.ones(1, 1, 10)
# 使用模型的 encode 方法对源序列进行编码，得到编码器的输出 memory
memory = test_model.encode(src, src_mask)
# 初始化目标序列 ys，开始时只包含一个起始标记（0）
ys = torch.zeros(1, 1).type_as(src)
# 循环 9 次，每次生成一个词
for i in range(9):
# 使用 decode 方法，传入编码器输出、源掩码、目标序列、未来信息掩码
out = test_model.decode(
memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
)
# 从解码器输出中取出最后一个位置的表示，通过 generator 生成概率分布
prob = test_model.generator(out[:, -1])
# 选择概率最大的词作为下一个词
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
# 将新生成的词添加到 ys 中，形成新的输入序列
ys = torch.cat(
[ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
)
# 打印最终生成的序列 ys
print("Example Untrained Model Prediction:", ys)
# 重复执行测试
def run_tests():
for _ in range(10):
inference_test()运行结果：
plaintext
Example Untrained Model Prediction: tensor([[ 0, 3, 10, 10, 10, 10, 10, 10, 3, 5]])
Example Untrained Model Prediction: tensor([[0, 2, 6, 2, 6, 2, 6, 2, 6, 2]])
Example Untrained Model Prediction: tensor([[0, 5, 3, 2, 2, 2, 2, 2, 2, 2]])
Example Untrained Model Prediction: tensor([[ 0, 10, 7, 10, 7, 10, 7, 10, 7, 10]])
Example Untrained Model Prediction: tensor([[ 0, 3, 10, 5, 10, 5, 10, 5, 10, 5]])
Example Untrained Model Prediction: tensor([[0, 7, 2, 4, 7, 2, 0, 7, 2, 3]])
Example Untrained Model Prediction: tensor([[ 0, 1, 2, 3, 4, 2, 3, 4, 10, 2]])
Example Untrained Model Prediction: tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
Example Untrained Model Prediction: tensor([[ 0, 10, 7, 4, 6, 10, 7, 4, 4, 4]])
Example Untrained Model Prediction: tensor([[0, 3, 6, 9, 3, 6, 9, 3, 6, 9]])