手撕 Transformer (5)：模型构建

前置知识：嵌入层和位置编码、编码器的实现、解码器和输出部分的实现。
之前的文章已经把构建 Transformer 所需的所有组件构建完了，这篇文章开始构建整个编码器-解码器结构。

1 编码器-解码器的代码实现

1. class EncoderDecoder(nn.Module):
2.     def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
3.         """
4.         encoder: 编码器模块，用于处理源序列
5.         decoder: 解码器模块，用于生成目标序列
6.         src_embed: 源序列嵌入层，将源序列tokens转换为向量表示
7.         tgt_embed: 目标序列嵌入层，将目标序列tokens转换为向量表示
8.         generator: 生成器模块，将解码器输出转换为目标词汇表上的概率分布
9.         """
10.         super(EncoderDecoder, self).__init__()
11.         self.encoder = encoder
12.         self.decoder = decoder
13.         self.src_embed = src_embed
14.         self.tgt_embed = tgt_embed
15.         self.generator = generator
17.     def forward(self, src, tgt, src_mask, tgt_mask):
18.         """
19.         Args:
20.             src: 源序列，形状为 [batch_size, src_seq_len]
21.             tgt: 目标序列，形状为 [batch_size, tgt_seq_len]
22.             src_mask: 源序列的掩码，形状为 [batch_size, 1, src_seq_len]，用于屏蔽填充位置
23.             tgt_mask: 目标序列的掩码，形状为 [batch_size, tgt_seq_len, tgt_seq_len]，用于屏蔽填充位置和未来位置
24.         Return:
25.             解码器的输出，形状为 [batch_size, tgt_seq_len, d_model]
26.         """
27.         # 先对源序列进行编码，得到编码器输出
28.         # 然后将编码器输出、源序列掩码、目标序列和目标序列掩码传递给解码器
29.         return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
31.     def encode(self, src, src_mask):
32.         """
33.         Args:
34.             src: 源序列，形状为 [batch_size, src_seq_len]
35.             src_mask: 源序列的掩码，形状为 [batch_size, 1, src_seq_len]，用于屏蔽填充位置
36.         Return:
37.             编码器的输出，形状为 [batch_size, src_seq_len, d_model]，作为解码器的输入
38.         """
39.         # (1) 对源序列进行嵌入处理，将tokens转换为向量表示
40.         # (2) 将嵌入结果和源序列掩码传递给编码器
41.         # (3) 返回编码器的输出，作为解码器的“记忆”输入
42.         return self.encoder(self.src_embed(src), src_mask)
44.     def decode(self, memory, src_mask, tgt, tgt_mask):
45.         """
46.         Args:
47.             memory: 编码器的输出
48.             src_mask: 源序列的掩码
49.             tgt: 目标序列
50.             tgt_mask: 目标序列的掩码
51.         """
52.         # (1) 对目标序列进行嵌入处理，将tokens转换为向量表示
53.         # (2) 将嵌入结果、编码器输出、源序列掩码、目标序列掩码传给解码器
54.         # (3) 返回解码器的输出
55.         return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

复制代码

实例化验证一下代码实现

1. vocab_size = 1000
2. d_model = 512
3. encoder = en
4. decoder = de
5. source_embed = nn.Embedding(vocab_size, d_model)
6. target_embed = nn.Embedding(vocab_size, d_model)
7. generator = gen
9. # 假设源数据与目标数据相同，实际中并不相同
10. source = target = Variable(torch.LongTensor([[100, 2, 421, 508], [491, 998, 1, 221]]))
12. # 假设src_mask与tgt_mask相同，实际中并不相同
13. source_mask = target_mask = Variable(torch.zeros(2, 4, 4))
15. ed = EncoderDecoder(encoder, decoder, source_embed, target_embed, generator)
16. ed_result = ed(source, target, source_mask, target_mask)
17. print(ed_result)
18. print(ed_result.shape)

复制代码

运行结果

1. tensor([[[-0.6909, -0.9724,  0.6361,  ...,  0.0525,  0.5336, -0.4723],
2.          [-1.7389, -1.2510,  0.7788,  ..., -0.0156,  0.0258, -0.9280],
3.          [-1.7369, -1.2625,  0.7260,  ..., -0.4470, -1.0672,  0.1064],
4.          [-1.3041, -1.3951, -0.0759,  ..., -0.1239, -0.6066, -0.8090]],
6.         [[-0.0288,  0.2606,  1.1938,  ...,  0.3468, -0.5532, -1.2409],
7.          [-0.2443, -0.1397,  0.7443,  ..., -0.2610,  0.6653, -1.7663],
8.          [-0.9244, -0.9061,  1.7313,  ...,  1.1313,  1.1599, -2.4360],
9.          [ 0.0872, -0.6195,  0.7045,  ..., -0.0286,  0.4245, -0.6196]]],
10.        grad_fn=)
11. torch.Size([2, 4, 512])

复制代码

2 完整模型构建

基于上述编码器-解码器结构构建用于训练的、完整的 Transformer 模型。

1. def make_model(
2.     src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
3. ):
4.     """
5.     Args:
6.         src_vocab: 源语言词汇表大小
7.         tgt_vocab: 目标语言词汇表大小
8.         N: 编码器和解码器的层数
9.         d_model: 模型的隐藏维度
10.         d_ff: 前馈网络的隐藏层维度
11.         h: 多头注意力的头数
12.         dropout: Dropout 丢弃率
13.     """
14.     # 创建一个深拷贝函数，用于复制组件实例，避免共享参数
15.     c = copy.deepcopy
16.     # 初始化多头注意力机制，处理不同子空间的注意力计算
17.     attn = MultiHeadedAttention(h, d_model)
18.     # 初始化位置前馈网络，对每个位置的表示进行非线性变换
19.     ff = PositionwiseFeedForward(d_model, d_ff, dropout)
20.     # 初始化位置编码，为输入序列添加位置信息
21.     position = PositionalEncoding(d_model, dropout)
22.     # 构建完整模型
23.     model = EncoderDecoder(
24.         # 编码器由 N 个 EncoderLayer 组成，每个层包含自注意力机制(attn)和前馈网络(ff)
25.         Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
26.         # 解码器由 N 个 DecoderLayer 组成，每个层包含自注意力机制(attn)、编码器-解码器注意力机制(attn)和前馈网络(ff)
27.         Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
28.         # 源嵌入：将源语言词汇转换为词向量，并添加位置编码
29.         nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
30.         # 目标嵌入：将目标语言词汇转换为词向量，并添加位置编码
31.         nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
32.         # 生成器：将解码器输出转换为目标词汇表的概率分布
33.         Generator(d_model, tgt_vocab),
34.     )
36.     # 参数初始化
37.     # 使用 Xavier 均匀分布初始化所有维度大于 1 的参数，
38.     # 这有助于模型的稳定训练和收敛
39.     for p in model.parameters():
40.         if p.dim() > 1:
41.             nn.init.xavier_uniform_(p)
42.     return model

复制代码

验证一下完整模型的代码实现

1.     source_vocab = 11
2.     target_vocab = 11
3.     N = 6
4.     res = make_model(source_vocab, target_vocab, N)
5.     print(res)

复制代码

运行结果

1. EncoderDecoder(
2.   (encoder): Encoder(
3.     (layers): ModuleList(
4.       (0-5): 6 x EncoderLayer(
5.         (self_attn): MultiHeadedAttention(
6.           (linears): ModuleList(
7.             (0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
8.           )
9.           (dropout): Dropout(p=0.1, inplace=False)
10.         )
11.         (feed_forward): PositionwiseFeedForward(
12.           (w_1): Linear(in_features=512, out_features=2048, bias=True)
13.           (w_2): Linear(in_features=2048, out_features=512, bias=True)
14.           (dropout): Dropout(p=0.1, inplace=False)
15.         )
16.         (sublayer): ModuleList(
17.           (0-1): 2 x SublayerConnection(
18.             (norm): LayerNorm()
19.             (dropout): Dropout(p=0.1, inplace=False)
20.           )
21.         )
22.       )
23.     )
24.     (norm): LayerNorm()
25.   )
26.   (decoder): Decoder(
27.     (layers): ModuleList(
28.       (0-5): 6 x DecoderLayer(
29.         (self_attn): MultiHeadedAttention(
30.           (linears): ModuleList(
31.             (0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
32.           )
33.           (dropout): Dropout(p=0.1, inplace=False)
34.         )
35.         (src_attn): MultiHeadedAttention(
36.           (linears): ModuleList(
37.             (0-3): 4 x Linear(in_features=512, out_features=512, bias=True)
38.           )
39.           (dropout): Dropout(p=0.1, inplace=False)
40.         )
41.         (feed_forward): PositionwiseFeedForward(
42.           (w_1): Linear(in_features=512, out_features=2048, bias=True)
43.           (w_2): Linear(in_features=2048, out_features=512, bias=True)
44.           (dropout): Dropout(p=0.1, inplace=False)
45.         )
46.         (sublayer): ModuleList(
47.           (0-2): 3 x SublayerConnection(
48.             (norm): LayerNorm()
49.             (dropout): Dropout(p=0.1, inplace=False)
50.           )
51.         )
52.       )
53.     )
54.     (norm): LayerNorm()
55.   )
56.   (src_embed): Sequential(
57.     (0): Embeddings(
58.       (lut): Embedding(11, 512)
59.     )
60.     (1): PositionalEncoding(
61.       (dropout): Dropout(p=0.1, inplace=False)
62.     )
63.   )
64.   (tgt_embed): Sequential(
65.     (0): Embeddings(
66.       (lut): Embedding(11, 512)
67.     )
68.     (1): PositionalEncoding(
69.       (dropout): Dropout(p=0.1, inplace=False)
70.     )
71.   )
72.   (generator): Generator(
73.     (proj): Linear(in_features=512, out_features=11, bias=True)
74.   )
75. )

复制代码

3 推理测试的代码实现

1. def inference_test():
2.     # 创建一个小型 Transformer 模型，词汇表大小为 11，编码器和解码器各 2 层
3.     test_model = make_model(11, 11, 2)
4.     # 将模型设置为评估模式
5.     test_model.eval()
6.     # 创建一个长度为 10 的源序列
7.     src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
8.     # 创建一个全 1 的源掩码，表示所有位置都是有效的（非填充）
9.     src_mask = torch.ones(1, 1, 10)
11.     # 使用模型的 encode 方法对源序列进行编码，得到编码器的输出 memory
12.     memory = test_model.encode(src, src_mask)
13.     # 初始化目标序列 ys，开始时只包含一个起始标记（0）
14.     ys = torch.zeros(1, 1).type_as(src)
16.     # 循环 9 次，每次生成一个词
17.     for i in range(9):
18.         # 使用 decode 方法，传入编码器输出、源掩码、目标序列、未来信息掩码
19.         out = test_model.decode(
20.             memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
21.         )
22.         # 从解码器输出中取出最后一个位置的表示，通过 generator 生成概率分布
23.         prob = test_model.generator(out[:, -1])
24.         # 选择概率最大的词作为下一个词
25.         _, next_word = torch.max(prob, dim=1)
26.         next_word = next_word.data[0]
27.         # 将新生成的词添加到 ys 中，形成新的输入序列
28.         ys = torch.cat(
29.             [ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
30.         )
32.     # 打印最终生成的序列 ys
33.     print("Example Untrained Model Prediction:", ys)
35. # 重复执行测试
36. def run_tests():
37.     for _ in range(10):
38.         inference_test()

复制代码

运行结果：

1. Example Untrained Model Prediction: tensor([[ 0,  3, 10, 10, 10, 10, 10, 10,  3,  5]])
2. Example Untrained Model Prediction: tensor([[0, 2, 6, 2, 6, 2, 6, 2, 6, 2]])
3. Example Untrained Model Prediction: tensor([[0, 5, 3, 2, 2, 2, 2, 2, 2, 2]])
4. Example Untrained Model Prediction: tensor([[ 0, 10,  7, 10,  7, 10,  7, 10,  7, 10]])
5. Example Untrained Model Prediction: tensor([[ 0,  3, 10,  5, 10,  5, 10,  5, 10,  5]])
6. Example Untrained Model Prediction: tensor([[0, 7, 2, 4, 7, 2, 0, 7, 2, 3]])
7. Example Untrained Model Prediction: tensor([[ 0,  1,  2,  3,  4,  2,  3,  4, 10,  2]])
8. Example Untrained Model Prediction: tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
9. Example Untrained Model Prediction: tensor([[ 0, 10,  7,  4,  6, 10,  7,  4,  4,  4]])
10. Example Untrained Model Prediction: tensor([[0, 3, 6, 9, 3, 6, 9, 3, 6, 9]])

复制代码

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！