自注意力同时具有并行计算和最短的最大路径长度这两个优势。Transformer 模型完全基于注意力机制，没有任何卷积层或循环神经网络层。尽管 Transformer 最初是应用于在文本数据上的序列到序列学习，但现在已经推广到各种现代的深度学习中，例如语言、视觉、语音和强化学习领域。

10.7.1 模型

Transformer 作为编码器－解码器架构的一个实例，其编码器和解码器是基于自注意力的模块叠加而成的，源（输入）序列和目标（输出）序列的嵌入（embedding）表示将加上位置编码（positional encoding），再分别输入到编码器和解码器中。

结构简介：

编码器：
- 由多个相同的层叠加而成，每个层有 两个子层（sublayer）：
  - 第一个子层为 多头自注意力（multi-head self-attention）汇聚
  - 第二个子层为 基于位置的前馈网络（positionwise feed-forward network）
- 在计算编码器的自注意力时，查询、键和值都来自前一个编码器层的输出
- 每个子层都采用了残差连接（residual connection）
- 残差连接的加法计算之后紧接着应用层规范化（layer normalization）
解码器：
- 由多个相同的层叠加而成，除了编码器中描述的两个子层之外，解码器还在这两个子层之间插入了第三个子层：
  - 编码器－解码器注意力（encoder-decoder attention）层
- 在解码器自注意力中，查询、键和值都来自上一个解码器层的输出
- 在第三个子层编码器－解码器注意力层中，查询来自前一个解码器层的输出，而键和值来自整个编码器的输出

import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

10.7.2 基于位置的前馈网络

基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的（positionwise）的原因。

名字很帅，其实就是全连接，但隐藏层的 MLP : )

输入X的形状（批量大小，时间步数或序列长度，隐单元数或特征维度）将被一个两层的感知机转换成形状为（批量大小，时间步数，ffn_num_outputs）的输出张量。

#@save
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,
                 **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))

ffn = PositionWiseFFN(4, 4, 8)
ffn.eval()
ffn(torch.ones((2, 3, 4)))[0]  # 把最后一个维度升上去

tensor([[-0.2199,  0.1357, -0.2216,  0.1659,  0.5388, -0.4541,  0.2121, -0.1025],
        [-0.2199,  0.1357, -0.2216,  0.1659,  0.5388, -0.4541,  0.2121, -0.1025],
        [-0.2199,  0.1357, -0.2216,  0.1659,  0.5388, -0.4541,  0.2121, -0.1025]],
       grad_fn=<SelectBackward0>)

10.7.3 残差连接和层规范化

批量归一化对每个特征/通道里的元素进行归一化，不适合序列长度会变的 NLP 应用。

层规范化和批量规范化的目标相同，但层规范化是基于特征维度进行规范化，即对每个样本里的元素进行归一化。

ln = nn.LayerNorm(2)
bn = nn.BatchNorm1d(2)
X = torch.tensor([[1, 2], [2, 3]], dtype=torch.float32)
# 在训练模式下计算X的均值和方差
print('layer norm:', ln(X), '\nbatch norm:', bn(X))  # layer norm 规范化的是每个样本（行）  batch norm 规范化的是每个特征（列）

layer norm: tensor([[-1.0000,  1.0000],
        [-1.0000,  1.0000]], grad_fn=<NativeLayerNormBackward0>) 
batch norm: tensor([[-1.0000, -1.0000],
        [ 1.0000,  1.0000]], grad_fn=<NativeBatchNormBackward0>)

#@save
class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)  # 暂退法也被作为正则化方法使用
        self.ln = nn.LayerNorm(normalized_shape)  # 层规范化

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)  # 残差连接

add_norm = AddNorm([3, 4], 0.5)
add_norm.eval()
add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape  # 残差连接要求两个输入的形状相同

torch.Size([2, 3, 4])

10.7.4 编码器

EncoderBlock 类包含两个子层：多头自注意力和基于位置的前馈网络，这两个子层都使用了残差连接和紧随的层规范化。

#@save
class EncoderBlock(nn.Module):
    """Transformer编码器块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,  # 写的很多，实际都是一个数
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.attention = d2l.MultiHeadAttention(  # 第一层 多头注意力层
            key_size, query_size, value_size, num_hiddens, num_heads, dropout,
            use_bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)  # 第一个层规范化
        self.ffn = PositionWiseFFN(  # 第二层 前馈网络层
            ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(norm_shape, dropout)  # 第二个规范化层

    def forward(self, X, valid_lens):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))  # 第一层
        return self.addnorm2(Y, self.ffn(Y))  # 第二层

X = torch.ones((2, 100, 24))
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
encoder_blk.eval()
encoder_blk(X, valid_lens).shape  # Transformer编码器中的任何层都不会改变其输入的形状

torch.Size([2, 100, 24])

以下对 num_layers个EncoderBlock 类的实例进行了堆叠。

这里使用的是值范围在 -1 和 1 之间的固定位置编码，因此通过学习得到的输入的嵌入表示的值需要先乘以嵌入维度的平方根进行重新缩放，然后再与位置编码相加。

#@save
class TransformerEncoder(d2l.Encoder):
    """Transformer编码器"""
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, use_bias=False, **kwargs):
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)  # 嵌入层
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)  # 位置编码
        self.blks = nn.Sequential()  # 编码器
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),  # 添加编码器的各层
                EncoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, use_bias))

    def forward(self, X, valid_lens, *args):
        # 因为位置编码值在-1和1之间，数值比较小，因此嵌入值乘以嵌入维度的平方根缩放到差不多大小，然后再与位置编码相加。
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)  # 存注意力汇聚权重用的
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)  # 一层一层的丢进去进行 attention
            self.attention_weights[  # 记录注意力汇聚权重
                i] = blk.attention.attention.attention_weights
        return X

encoder = TransformerEncoder(  # 创建一个两层的Transformer编码器
    200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
encoder.eval()
encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape  # 输出的形状是（批量大小，时间步数目，num_hiddens）

torch.Size([2, 100, 24])

10.7.5 解码器

DecoderBlock 类中实现的每个层包含了三个子层：解码器自注意力、“编码器-解码器”注意力和基于位置的前馈网络。这些子层也都被残差连接和紧随的层规范化围绕。

class DecoderBlock(nn.Module):
    """解码器中第i个块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        self.attention1 = d2l.MultiHeadAttention(  # 第一层 解码器的自注意力层
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)  # 第一个层规范化
        self.attention2 = d2l.MultiHeadAttention(  # 第二层 “编-解”注意力层
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)  # 第二个层规范化
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,  # 第三层 掐灭亏网络层
                                   num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)  # 第三个层规范化

    def forward(self, X, state):
        enc_outputs, enc_valid_lens = state[0], state[1]
        if state[2][self.i] is None:  # 训练阶段，输出序列的所有词元都在同一时间处理，因此state[2][self.i]初始化为None。
            key_values = X
        else:  # 预测阶段，输出序列是通过词元一个接着一个解码的，因此需要把直到当前时间步第i个块解码的输出在state[2][self.i]里面存着
            key_values = torch.cat((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values
        if self.training:
            batch_size, num_steps, _ = X.shape
            # dec_valid_lens 的开头:(batch_size,num_steps)，其中每一行是[1,2,...,num_steps]
            dec_valid_lens = torch.arange(  # 只要训练状态下需要遮掉后面的内容
                1, num_steps + 1, device=X.device).repeat(batch_size, 1)
        else:  # 预测模式后面空白的 不用管
            dec_valid_lens = None

      
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)  # 自注意力
        Y = self.addnorm1(X, X2)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)  # 编码器－解码器注意力。enc_outputs的开头:(batch_size,num_steps,num_hiddens)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, 0)
decoder_blk.eval()
X = torch.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
decoder_blk(X, state)[0].shape  # 过一遍形状不变的

torch.Size([2, 100, 24])

class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)  # 词嵌入层
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)  # 位置编码
        self.blks = nn.Sequential()  # 解码器
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),  # 添加解码器的各层
                DecoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, i))
        self.dense = nn.Linear(num_hiddens, vocab_size)  # 最后的全连接层

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]  # 最后那个是预测时存东西用的

    def forward(self, X, state):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))  # 叠加位置编码
        self._attention_weights = [[None] * len(self.blks) for _ in range (2)]  # 存注意力汇聚权重用的
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            self._attention_weights[0][  # 存解码器自注意力权重
                i] = blk.attention1.attention.attention_weights
            self._attention_weights[1][  # 存“编码器－解码器”自注意力权重
                i] = blk.attention2.attention.attention_weights
        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights

10.7.6 训练

依照 Transformer 架构来实例化编码器－解码器模型。

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()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4  # 都使用4头注意力
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]

train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

encoder = TransformerEncoder(  # 整一个编码器
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
decoder = TransformerDecoder(  # 整一个解码器
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)  # 组网
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)  # 并行度应该挺高的 训起来不慢 1分9秒

loss 0.031, 6866.6 tokens/sec on cuda:0

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

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')

go . => va !,  bleu 1.000
i lost . => j'ai perdu .,  bleu 1.000
he's calm . => il est paresseux .,  bleu 0.658
i'm home . => je suis chez moi .,  bleu 1.000

可视化 Transformer 的注意力权重。

enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
    -1, num_steps))
enc_attention_weights.shape  # 编码器自注意力权重的形状为（编码器层数，注意力头数，num_steps或查询的数目，num_steps或“键－值”对的数目）

torch.Size([2, 4, 10, 10])

d2l.show_heatmaps(  # 逐行呈现编码器的两层多头注意力的权重
    enc_attention_weights.cpu(), xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))  # 可以看到每个注意力头的注意力都不大一样。

用零填充被掩蔽住的注意力权重后，可视化解码器的自注意力权重和“编码器－解码器”的注意力权重。

解码器的自注意力权重和“编码器－解码器”的注意力权重都有相同的查询：即以序列开始词元（beginning-of-sequence,BOS）打头，再与后续输出的词元共同组成序列。

dec_attention_weights_2d = [head[0].tolist()
                            for step in dec_attention_weight_seq
                            for attn in step for blk in attn for head in blk]
dec_attention_weights_filled = torch.tensor(
    pd.DataFrame(dec_attention_weights_2d).fillna(0.0).values)  # 用零填充被掩蔽住的注意力权重
dec_attention_weights = dec_attention_weights_filled.reshape((-1, 2, num_layers, num_heads, num_steps))
dec_self_attention_weights, dec_inter_attention_weights = \
    dec_attention_weights.permute(1, 2, 3, 0, 4)
dec_self_attention_weights.shape, dec_inter_attention_weights.shape

(torch.Size([2, 4, 6, 10]), torch.Size([2, 4, 6, 10]))

由于解码器自注意力的自回归属性，查询不会对当前位置之后的“键－值”对进行注意力计算。

d2l.show_heatmaps(  # 逐行呈现解码器的多头自注意力的权重
    dec_self_attention_weights[:, :, :, :len(translation.split()) + 1],
    xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(7, 3.5))

与编码器的自注意力的情况类似，通过指定输入序列的有效长度，输出序列的查询不会与输入序列中填充位置的词元进行注意力计算。

d2l.show_heatmaps(  # 逐行呈现解码器的编-解多头自注意力的权重
    dec_inter_attention_weights, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))

练习

（1）在实验中训练更深的 Transformer 将如何影响训练速度和翻译效果？

更慢了，注意力越往后越浓重。

num_hiddens, num_layers_deeper, dropout, batch_size, num_steps = 32, 4, 0.1, 64, 10  # 加深到 4 层
lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4  # 还是4头注意力
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]

train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

encoder_deeper = TransformerEncoder(
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers_deeper, dropout)
decoder_deeper = TransformerDecoder(
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers_deeper, dropout)
net_deeper = d2l.EncoderDecoder(encoder_deeper, decoder_deeper)
d2l.train_seq2seq(net_deeper, train_iter, lr, num_epochs, tgt_vocab, device)  # 时间慢到2分了 怎么精度跌了

loss 0.063, 4044.0 tokens/sec on cuda:0

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net_deeper, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')

go . => va le chercher !,  bleu 0.000
i lost . => je me suis tombé .,  bleu 0.000
he's calm . => il est malade .,  bleu 0.658
i'm home . => je suis sûr .,  bleu 0.512

enc_attention_weights_deeper = torch.cat(net_deeper.encoder.attention_weights, 0).reshape((num_layers_deeper, num_heads,
    -1, num_steps))

d2l.show_heatmaps(
    enc_attention_weights_deeper.cpu(), xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(10, 8))

dec_attention_weights_2d_deeper = [head[0].tolist()
                            for step in dec_attention_weight_seq
                            for attn in step for blk in attn for head in blk]
dec_attention_weights_filled_deeper = torch.tensor(
    pd.DataFrame(dec_attention_weights_2d_deeper).fillna(0.0).values)  # 用零填充被掩蔽住的注意力权重
dec_attention_weights_deeper = dec_attention_weights_filled_deeper.reshape((-1, 2, num_layers_deeper, num_heads, num_steps))
dec_self_attention_weights_deeper, dec_inter_attention_weights_deeper = \
    dec_attention_weights_deeper.permute(1, 2, 3, 0, 4)

d2l.show_heatmaps(  # 逐行呈现解码器的多头自注意力的权重
    dec_self_attention_weights_deeper[:, :, :, :len(translation.split()) + 1],
    xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(10, 8))

d2l.show_heatmaps(
    dec_inter_attention_weights_deeper, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(10, 8))