1. import math
  2. import pandas as pd
  3. import torch
  4. from torch import nn
  5. from d2l import torch as d2l

10.7.2. 基于位置的前馈网络

  1. class PositionWiseFFN(nn.Module):
  2.     def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,
  3.                  **kwargs):
  4.         super(PositionWiseFFN, self).__init__(**kwargs)
  5.         self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
  6.         self.relu = nn.ReLU()
  7.         self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)
  9.     def forward(self, X):
  10.         return self.dense2(self.relu(self.dense1(X)))
  1. ffn = PositionWiseFFN(4, 4, 8)
  2. ffn.eval()
  3. ffn(torch.ones((2, 3, 4)))[0]
  1. tensor([[ 0.0860,  0.1906,  0.1094, -0.7480, -0.6123,  0.2456, -0.3660, -0.5613],
  2.         [ 0.0860,  0.1906,  0.1094, -0.7480, -0.6123,  0.2456, -0.3660, -0.5613],
  3.         [ 0.0860,  0.1906,  0.1094, -0.7480, -0.6123,  0.2456, -0.3660, -0.5613]],
  4.        grad_fn=<SelectBackward>)

10.7.3. 残差连接和层归一化

  1. ln = nn.LayerNorm(2)
  2. bn = nn.BatchNorm1d(2)
  3. X = torch.tensor([[1, 2], [2, 3]], dtype=torch.float32)
  4. # 在训练模式下计算 `X` 的均值和方差
  5. print('layer norm:', ln(X), '\nbatch norm:', bn(X))
  1. layer norm: tensor([[-1.0000,  1.0000],
  2.         [-1.0000,  1.0000]], grad_fn=<NativeLayerNormBackward>)
  3. batch norm: tensor([[-1.0000, -1.0000],
  4.         [ 1.0000,  1.0000]], grad_fn=<NativeBatchNormBackward>)
  1. class AddNorm(nn.Module):
  2.     def __init__(self, normalized_shape, dropout, **kwargs):
  3.         super(AddNorm, self).__init__(**kwargs)
  4.         self.dropout = nn.Dropout(dropout)
  5.         self.ln = nn.LayerNorm(normalized_shape)
  7.     def forward(self, X, Y):
  8.         return self.ln(self.dropout(Y) + X)
  1. add_norm = AddNorm([3, 4], 0.5) # Normalized_shape is input.size()[1:]
  2. add_norm.eval()
  3. add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape
  1. torch.Size([2, 3, 4])

10.7.4. 编码器

  1. class EncoderBlock(nn.Module):
  2.     def __init__(self, key_size, query_size, value_size, num_hiddens,
  3.                  norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
  4.                  dropout, use_bias=False, **kwargs):
  5.         super(EncoderBlock, self).__init__(**kwargs)
  6.         self.attention = d2l.MultiHeadAttention(
  7.             key_size, query_size, value_size, num_hiddens, num_heads, dropout,
  8.             use_bias)
  9.         self.addnorm1 = AddNorm(norm_shape, dropout)
  10.         self.ffn = PositionWiseFFN(
  11.             ffn_num_input, ffn_num_hiddens, num_hiddens)
  12.         self.addnorm2 = AddNorm(norm_shape, dropout)
  14.     def forward(self, X, valid_lens):
  15.         Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
  16.         return self.addnorm2(Y, self.ffn(Y))
  1. X = torch.ones((2, 100, 24))
  2. valid_lens = torch.tensor([3, 2])
  3. encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
  4. encoder_blk.eval()
  5. encoder_blk(X, valid_lens).shape
  1. torch.Size([2, 100, 24])
  1. class TransformerEncoder(d2l.Encoder):
  2.     def __init__(self, vocab_size, key_size, query_size, value_size,
  3.                  num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
  4.                  num_heads, num_layers, dropout, use_bias=False, **kwargs):
  5.         super(TransformerEncoder, self).__init__(**kwargs)
  6.         self.num_hiddens = num_hiddens
  7.         self.embedding = nn.Embedding(vocab_size, num_hiddens)
  8.         self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
  9.         self.blks = nn.Sequential()
  10.         for i in range(num_layers):
  11.             self.blks.add_module("block"+str(i),
  12.                 EncoderBlock(key_size, query_size, value_size, num_hiddens,
  13.                              norm_shape, ffn_num_input, ffn_num_hiddens,
  14.                              num_heads, dropout, use_bias))
  16.     def forward(self, X, valid_lens, *args):
  17.         # 因为位置编码值在 -1 和 1 之间，
  18.         # 因此嵌入值乘以嵌入维度的平方根进行缩放，
  19.         # 然后再与位置编码相加。
  20.         X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
  21.         self.attention_weights = [None] * len(self.blks)
  22.         for i, blk in enumerate(self.blks):
  23.             X = blk(X, valid_lens)
  24.             self.attention_weights[
  25.                 i] = blk.attention.attention.attention_weights
  26.         return X
  1. encoder = TransformerEncoder(
  2.     200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
  3. encoder.eval()
  4. encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape
  1. torch.Size([2, 100, 24])

10.7.5. 解码器

  1. class DecoderBlock(nn.Module):
  2.     """解码器中第 i 个块"""
  3.     def __init__(self, key_size, query_size, value_size, num_hiddens,
  4.                  norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
  5.                  dropout, i, **kwargs):
  6.         super(DecoderBlock, self).__init__(**kwargs)
  7.         self.i = i
  8.         self.attention1 = d2l.MultiHeadAttention(
  9.             key_size, query_size, value_size, num_hiddens, num_heads, dropout)
  10.         self.addnorm1 = AddNorm(norm_shape, dropout)
  11.         self.attention2 = d2l.MultiHeadAttention(
  12.             key_size, query_size, value_size, num_hiddens, num_heads, dropout)
  13.         self.addnorm2 = AddNorm(norm_shape, dropout)
  14.         self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,
  15.                                    num_hiddens)
  16.         self.addnorm3 = AddNorm(norm_shape, dropout)
  18.     def forward(self, X, state):
  19.         enc_outputs, enc_valid_lens = state[0], state[1]
  20.         # 训练阶段，输出序列的所有词元都在同一时间处理，
  21.         # 因此 `state[2][self.i]` 初始化为 `None`。
  22.         # 预测阶段，输出序列是通过词元一个接着一个解码的，
  23.         # 因此 `state[2][self.i]` 包含着直到当前时间步第 `i` 个块解码的输出表示
  24.         if state[2][self.i] is None:
  25.             key_values = X
  26.         else:
  27.             key_values = torch.cat((state[2][self.i], X), axis=1)
  28.         state[2][self.i] = key_values
  29.         if self.training:
  30.             batch_size, num_steps, _ = X.shape
  31.             # `dec_valid_lens` 的开头: (`batch_size`, `num_steps`),
  32.             # 其中每一行是 [1, 2, ..., `num_steps`]
  33.             dec_valid_lens = torch.arange(
  34.                 1, num_steps + 1, device=X.device).repeat(batch_size, 1)
  35.         else:
  36.             dec_valid_lens = None
  38.         # 自注意力
  39.         X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
  40.         Y = self.addnorm1(X, X2)
  41.         # 编码器－解码器注意力。
  42.         # `enc_outputs` 的开头: (`batch_size`, `num_steps`, `num_hiddens`)
  43.         Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
  44.         Z = self.addnorm2(Y, Y2)
  45.         return self.addnorm3(Z, self.ffn(Z)), state
  1. decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, 0)
  2. decoder_blk.eval()
  3. X = torch.ones((2, 100, 24))
  4. state = [encoder_blk(X, valid_lens), valid_lens, [None]]
  5. decoder_blk(X, state)[0].shape
  1. class TransformerDecoder(d2l.AttentionDecoder):
  2.     def __init__(self, vocab_size, key_size, query_size, value_size,
  3.                  num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
  4.                  num_heads, num_layers, dropout, **kwargs):
  5.         super(TransformerDecoder, self).__init__(**kwargs)
  6.         self.num_hiddens = num_hiddens
  7.         self.num_layers = num_layers
  8.         self.embedding = nn.Embedding(vocab_size, num_hiddens)
  9.         self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
  10.         self.blks = nn.Sequential()
  11.         for i in range(num_layers):
  12.             self.blks.add_module("block"+str(i),
  13.                 DecoderBlock(key_size, query_size, value_size, num_hiddens,
  14.                              norm_shape, ffn_num_input, ffn_num_hiddens,
  15.                              num_heads, dropout, i))
  16.         self.dense = nn.Linear(num_hiddens, vocab_size)
  18.     def init_state(self, enc_outputs, enc_valid_lens, *args):
  19.         return [enc_outputs, enc_valid_lens, [None] * self.num_layers]
  21.     def forward(self, X, state):
  22.         X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
  23.         self._attention_weights = [[None] * len(self.blks) for _ in range (2)]
  24.         for i, blk in enumerate(self.blks):
  25.             X, state = blk(X, state)
  26.             # 解码器自注意力权重
  27.             self._attention_weights[0][
  28.                 i] = blk.attention1.attention.attention_weights
  29.             # “编码器－解码器”自注意力权重
  30.             self._attention_weights[1][
  31.                 i] = blk.attention2.attention.attention_weights
  32.         return self.dense(X), state
  34.     @property
  35.     def attention_weights(self):
  36.         return self._attention_weights

10.7.6. 训练

  1. num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
  2. lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
  3. ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
  4. key_size, query_size, value_size = 32, 32, 32
  5. norm_shape = [32]
  7. train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
  9. encoder = TransformerEncoder(
  10.     len(src_vocab), key_size, query_size, value_size, num_hiddens,
  11.     norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
  12.     num_layers, dropout)
  13. decoder = TransformerDecoder(
  14.     len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
  15.     norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
  16.     num_layers, dropout)
  17. net = d2l.EncoderDecoder(encoder, decoder)
  18. d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
  1. loss 0.030, 5244.8 tokens/sec on cuda:0

image.png
训练结束后，使用 Transformer 模型将一些英语句子翻译成法语，
并且计算它们的 BLEU 分数。

  1. engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
  2. fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
  3. for eng, fra in zip(engs, fras):
  4.     translation, dec_attention_weight_seq = d2l.predict_seq2seq(
  5.         net, eng, src_vocab, tgt_vocab, num_steps, device, True)
  6.     print(f'{eng} => {translation}, ',
  7.           f'bleu {d2l.bleu(translation, fra, k=2):.3f}')
  1. go . => va !,  bleu 1.000
  2. i lost . => j'ai perdu .,  bleu 1.000
  3. he's calm . => il est calme .,  bleu 1.000
  4. i'm home . => je suis chez moi .,  bleu 1.000
  1. enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
  2.     -1, num_steps))
  3. enc_attention_weights.shape
  1. torch.Size([2, 4, 10, 10])