跟着李沐老师学习深度学习（十四）

注意力机制（Attention）

引入

• 心理学角度
  • 动物需要在复杂环境下有效关注值得注意的点
  • 心理学框架：人类根据随意线索和不随意线索选择注意力

注意力机制

• 之前所涉及到的卷积、全连接、池化层都只考虑不随意线索
• 而注意力机制则显示的考虑随意线索
  • 随意线索 被称为查询（query）
  • 每个输入是一个**值（value）和不随意线索（key）**的对
  • 通过注意力池化层来有偏向性的选择某些输入

背景

（1）非参注意力池化层

• 给定数据（x_i, y_i）,i = 1,… ,n
• 平均池化是最简单的方案：
  
  x是query
• 更好的方案是60年代提出来的Nadaraya-Watson 核回归：

K： 核，衡量x和x_i距离的函数；整个函数意思就是，当来了一个新数据的时候，我该怎么看，应该找和我相近的数据拿出来，别的不管。

（2）Nadaraya-Watson 核回归

• 使用高斯核 :
• 那么：
  
  其实就是一个softmax；如果一个键 x_i 越是接近给定的查询 x ， 那么分配给这个键对应值y_i的注意力权重就会越大， 也就“获得了更多的注意力”。

（3）参数化的注意力机制

• 非参数的Nadaraya-Watson核回归具有**一致性（consistency）**的优点： 如果有足够的数据，此模型会收敛到最优结果。 尽管如此，我们还是可以轻松地将可学习的参数集成到注意力汇聚中。
• 在之前基础上引入可以学习的w
  

总结

• Nadaraya-Watson核回归是具有注意力机制的机器学习范例。
• Nadaraya-Watson核回归的注意力汇聚是对训练数据中输出的加权平均。从注意力的角度来看，分配给每个值的注意力权重取决于将值所对应的键和查询作为输入的函数。
• 注意力汇聚可以分为非参数型和带参数型。

代码实现

# Nadaraya-Watson 核回归
import torch
from torch import nn
from d2l import torch as d2l

# 生成数据集
n_train = 50 # 训练样本数
x_train, _ = torch.sort(torch.rand(n_train) * 5)

def f(x):
    return 2 * torch.sin(x) + x ** 0.8

y_train = f(x_train) + torch.normal(0.0, 0.5, (n_train,)) # 训练样本的输出
# torch.arange是 PyTorch 中的函数，用于创建具有等间距值的一维张量，类似于 Python 的range函数，但返回的是张量。
# 其函数定义为torch.arange(start=0, end, step=1, *, dtype=None, layout=torch.strided, device=None, requires_grad=False)
# 一个包含从 0 开始，以 0.1 为间隔，到小于 5 的所有数值的一维张量
x_test = torch.arange(0, 5, 0.1)
y_truth = f(x_test)
n_test = len(x_test)
n_test

# 绘制所有的训练样本（样本由圆圈表示）
def plot_kernel_reg(y_hat):
    d2l.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'],
             xlim=[0, 5], ylim=[-1, 5])
    d2l.plt.plot(x_train, y_train, 'o', alpha=0.5)

# 先用最简单的估计器来解决：使用平均汇聚来计算所有训练样本输出值的平均值
y_hat = torch.repeat_interleave(y_train.mean(), n_test)
plot_kernel_reg(y_hat)

# 平均汇聚忽略了输入 𝑥𝑖
# 使用非参数注意力汇聚
# X_repeat的形状:(n_test,n_train),
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))
# x_train包含着键。attention_weights的形状：(n_test,n_train),
# 每一行都包含着要在给定的每个查询的值（y_train）之间分配的注意力权重
attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)
# y_hat的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights, y_train)
plot_kernel_reg(y_hat)

d2l.show_heatmaps(attention_weights.unsqueeze(0).unsqueeze(0),
                  xlabel='Sorted training inputs',
                  ylabel='Sorted testing inputs')

# 带参数注意力汇聚
# 批量矩阵的乘法
# 假定两个张量的形状分别是 (𝑛,𝑎,𝑏) 和 (𝑛,𝑏,𝑐) ， 它们的批量矩阵乘法输出的形状为 (𝑛,𝑎,𝑐)
X = torch.ones((2, 1, 4))
Y = torch.ones((2, 4, 6))
torch.bmm(X, Y).shape

# 在注意力机制的背景中，可以使用小批量矩阵乘法来计算小批量数据中的加权平均值
weights = torch.ones((2, 10)) * 0.1
values = torch.arange(20.0).reshape((2, 10))
# torch.bmm 函数用于执行批量矩阵乘法。它要求输入的两个张量必须都是 3 维的，并且满足一定的维度匹配规则：
# 第一个张量的形状为 (batch_size, m, n) 。
# 第二个张量的形状为 (batch_size, n, p) 。
# weights.unsqueeze(1)：unsqueeze 方法用于在指定位置插入一个维度为 1 的新维度
# weights 变为（2，1，10）
# values.unsqueeze(-1) 会在 values 张量的最后一个维度位置插入一个新维度
# values 变成（2，10，1）
torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1)),torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1)).shape

# 定义模型
class NWKernelRegression(nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.w = nn.Parameter(torch.rand((1,), requires_grad=True))

    def forward(self, queries, keys, values):
        # queries和attention_weights的形状为(查询个数，“键－值”对个数)
        queries = queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1]))
        self.attention_weights = nn.functional.softmax(
            -((queries - keys) * self.w)**2 / 2, dim=1)
        # values的形状为(查询个数，“键－值”对个数)
        return torch.bmm(self.attention_weights.unsqueeze(1),
                         values.unsqueeze(-1)).reshape(-1)

# 训练： 将训练数据集 变换为 键和值
# X_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输入
X_tile = x_train.repeat((n_train, 1))
# Y_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输出
Y_tile = y_train.repeat((n_train, 1))
# keys的形状:('n_train'，'n_train'-1)
keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))
# values的形状:('n_train'，'n_train'-1)
values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))

# 训练带参数的注意力汇聚模型 时，使用平方损失函数和随机梯度下降。
net = NWKernelRegression()
loss = nn.MSELoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])

for epoch in range(5):
    trainer.zero_grad()
    l = loss(net(x_train, keys, values), y_train)
    l.sum().backward()
    trainer.step()
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

#  在尝试拟合带噪声的训练数据时， [预测结果绘制]的线不如之前非参数模型的平滑。

# keys的形状:(n_test，n_train)，每一行包含着相同的训练输入（例如，相同的键）
keys = x_train.repeat((n_test, 1))
# value的形状:(n_test，n_train)
values = y_train.repeat((n_test, 1))
y_hat = net(x_test, keys, values).unsqueeze(1).detach()
plot_kernel_reg(y_hat)

# 为什么新的模型更不平滑了呢？ 
# 下面看一下输出结果的绘制图： 与非参数的注意力汇聚模型相比， 带参数的模型加入可学习的参数后，曲线在注意力权重较大的区域变得更不平滑
d2l.show_heatmaps(net.attention_weights.unsqueeze(0).unsqueeze(0),
                  xlabel='Sorted training inputs',
                  ylabel='Sorted testing inputs')

结果：
（1）使用平均汇聚（预测为一条直线）

（2）非参数注意力汇聚

（3）带参数的注意力汇聚

注意力分数

• 回顾：
  

• 扩展到高维度
• 用数学语言描述，假设有一个査询 q 和 m个“键一值”对(k1,v1),…,(km, vm)，，注意力汇聚函数 f 就被表示成值的加权和：
  
  其中查询 q 和键 ki 的注意力权重（标量） 是通过注意力评分函数 α 将两个向量映射成标量， 再经过softmax运算得到的：
  

注意力评分函数 α 的设计

• 思路1：additive attention（加性注意力）

  • 可学参数：
    

  

  • 等价于将key和value合并起来后放入到一个隐藏大小为h 输出大小为 1 的单隐藏MLP。
  • 好处：query、value 可以是不等长的。

• 思路2：scaled dot-product attention（缩放点积注意力）

  • 如果query和key都是同样的长度q,k ，使用点积可以得到计算效率更高的评分函数， 但是点积操作要求查询和键具有相同的长度d。 假设查询和键的所有元素都是独立的随机变量， 并且都满足零均值和单位方差， 那么两个向量的点积的均值为0，方差为d。
  • 为确保无论向量长度如何， 点积的方差在不考虑向量长度的情况下仍然是1， 我们再将点积除以 根号下d
    
  • 向量化版本
    

总结

• 注意力分数是query和key的相似度，注意力权重是分数的softmax结果
• 两种常见的分数计算
• 将query和key合并起来进入一个单输出单隐藏层的MLP
• 直接将query和key做内积

自注意力和位置编码

自注意力

• 给定序列
  

• 自注意力池化层 将x_i 当作 key、value、query来对序列抽取特征得到 y_1, …, y_n。这里：
  
  

• 与CNN，RNN对比
  
  
  k: 窗口大小；n:序列大小；d：特征维度
  CNN做序列：将序列当作一个一维的输入，只有一个宽（图片有高和宽）来处理文本序列；计算复杂度
  RNN：并行度差，要算完前一个再继续算。

• 计算复杂度：

  • CNN：其中， k是卷积核的数量， n是序列长度， d是特征维度。卷积操作通过卷积核在输入上滑动来提取特征，计算量与卷积核数量、序列长度和特征维度的平方相关 。
  • RNN：在 RNN 中，每个时间步的计算依赖于前一个时间步的状态，计算复杂度与序列长度和特征维度的平方成正比。
  • 自注意力：自注意力机制会计算序列中每个元素与其他所有元素的注意力权重，因此计算复杂度与序列长度的平方和特征维度成正比。

• 并行度

  • CNN：卷积操作可以在不同位置上并行计算，因为每个卷积核在不同位置的计算是相互独立的，所以并行度与序列长度成正比。
  • RNN：RNN 的计算具有顺序性，每个时间步的计算依赖于前一个时间步的输出，很难进行并行计算，并行度较低，通常为常数级。
  • 自注意力：自注意力机制可以同时计算序列中所有元素的注意力权重，因此并行度与序列长度成正比，能够高效地利用并行计算资源。

• 最长路径：

  • CNN：卷积操作通过逐层卷积来扩大感受野，信息在网络中的传播路径长度与序列长度除以卷积核的数量相关。
  • RNN：由于 RNN 是顺序处理序列数据，信息需要从序列的起始位置逐步传递到末尾位置，最长路径与序列长度成正比。
  • 自注意力：自注意力机制可以直接捕捉序列中任意两个元素之间的关系，信息传递不需要经过多个中间步骤，因此最长路径为常数级。（适合处理比较长的文本）

位置编码

• 与CNN、RNN不同，自注意力并没有记录位置信息
• 位置编码将位置信息注入到输入里
• P的元素如下计算
  
• 效果：（位移和周期的变化）
  为什么用sin、cos函数表示位置信息？——可以表示相对位置信息
  • 周期性：函数的周期性契合序列位置的规律特征，有助于模型捕捉位置相关规律。
  • 位置区分：不同位置函数值不同，形成编码差异，让模型辨别位置。
  • 可解释性：作为经典数学函数，性质明确，便于理解其对模型的作用。
  • 相对位置表达：能有效表征元素相对位置关系，对处理序列元素关系重要。
  • 泛化性：在不同长度序列上表现良好，训练模式可推广，适应性强。

总结

• 自注意力池化层将 x_i 当做key，value，query来对序列抽取特征
• 完全并行、最长序列为1、但对长序列计算复杂度高
• 位置编码在输入中加入位置信息，使得自注意力能够记忆位置信息

Transformer

• 基于编码器-解码器 架构来处理序列对
• 跟使用注意力的 Seq2Seq不同，transformer是纯基于注意力。

到这里，感觉对 rnn 和 transformer 有些迷惑的地方，稍微总结下：
RNN 和基于注意力模型的不同
（1）结构原理：
RNN：具有循环结构，通过隐藏状态在时间步之间传递信息，处理序列数据时，当前时刻的输出依赖于之前时刻的隐藏状态和当前输入。比如在处理文本序列时，会依次处理每个单词，将前一个单词的信息融入到当前单词的处理中。
基于注意力模型：注意力机制不是一种特定的网络结构，而是一种让模型聚焦于输入序列中重要部分的方法。它通过计算输入元素之间的注意力权重，来决定每个元素在当前计算中的重要程度，从而更灵活地处理序列数据。
（2）长序列处理能力：
RNN：处理长序列时存在梯度消失或梯度爆炸问题，使得它难以有效捕捉长距离依赖关系，记忆能力有限。
基于注意力模型：可以直接建立序列中任意位置之间的依赖关系，不受距离限制，能够更有效地处理长距离依赖问题。
（3）并行处理能力：
RNN：计算具有顺序性，必须按顺序依次处理每个时间步，难以并行计算，训练效率较低。
基于注意力模型：如 Transformer 中的自注意力机制，不依赖于序列元素的顺序，可以并行计算所有元素的注意力权重，训练效率大大提高。

seq2seq 和 transformer的不同：
（1）模型架构：
seq2seq：是一种编码器-解码器结构，编码器将输入序列映射为固定维度的向量表示，解码器将这个向量映射为目标序列。常用 RNN 或其变体（如 LSTM、GRU）作为编码器和解码器。
transformer：同样有编码器和解码器两部分，但完全基于自注意力机制。编码器由多层自注意力机制和前馈神经网络组成，将输入序列编码为向量表示；解码器除了自注意力机制和前馈神经网络外，还有交叉注意力子层用于结合编码器输出信息。
（2）注意力机制：
seq2seq：通常在解码器部分使用注意力机制，来帮助模型在生成输出时关注输入序列的不同部分，但这种注意力机制与 transformer 的自注意力机制有所不同。
transformer：引入多头自注意力机制，允许模型同时在不同的表示子空间上捕捉信息，能并行处理所有词汇之间的关系，更好地解决长距离依赖问题，增强了模型表达能力。

• transformer框架：
  • Transformer的编码器是由多个相同的层叠加而成的，每个层都有两个子层（子层表示为sublayer）。第一个子层是多头自注意力（multi-head self-attention）汇聚；第二个子层是基于位置的前馈网络（positionwise feed-forward network） 。
    
    图片详解：
• 编码器（Encoder）
  • 嵌入层（Embedding Layer）：接收源序列（如源语言句子）作为输入，将每个词转换为对应的词嵌入向量，赋予词的语义信息。
  • 位置编码（Positional Encoding）：由于 Transformer 模型本身不具备对序列顺序的感知能力，位置编码通过特定的方式（如使用 sin 和 cos 函数）为每个位置生成独特的编码，然后将其与词嵌入向量相加，从而为模型引入序列的顺序信息。
  • Transformer 块（Transformer Block）：由多个（图中为 n 个）相同的模块堆叠而成，每个模块包含两个主要子层：
    • 多头注意力（Multi - Head Attention）：通过多个头并行计算注意力，捕捉输入序列中不同位置之间的依赖关系，从不同角度对输入进行建模，获取丰富的语义信息。
    • 逐位前馈网络（Feed - Forward Network）：对多头注意力的输出进行进一步处理，通过线性变换和非线性激活函数，增强模型对特征的表示能力。
  • 加 & 规范化（Add & Normalization）：这是一种残差连接和层归一化的操作。残差连接将输入直接加到子层的输出上，有助于解决深层网络训练中的梯度消失等问题；层归一化则对数据进行归一化处理，使数据分布更稳定，加速模型收敛。
• 解码器（Decoder）
  • 嵌入层（Embedding Layer）：与编码器类似，接收目标序列（如目标语言句子）作为输入，将词转换为词嵌入向量。
  • 位置编码（Positional Encoding）：同样为目标序列添加位置信息，使其能感知序列顺序。
  • Transformer 块（Transformer Block）：由多个（图中为 n 个）模块堆叠而成，每个模块包含三个子层：
    • 有掩码的多头注意力（Masked Multi-Head Attention）：在生成目标序列时，为了防止解码器提前看到未来位置的信息，使用掩码机制对后续位置进行屏蔽，确保在生成当前位置时只能依赖于已经生成的位置信息。
    • 多头注意力（Multi - Head Attention）：用于接收编码器的输出，并计算与解码器输入之间的注意力，使解码器能够利用编码器对源序列的编码信息，指导目标序列的生成。
    • 逐位前馈网络（Feed - Forward Network）：功能与编码器中的前馈网络相同，对经过注意力机制处理后的信息进一步加工。
  • 加 & 规范化（Add & Normalization）：在每个子层后都有此操作，作用与编码器中的一致，保证信息传递的稳定性和有效性。
  • 全连接层（Fully - Connected Layer）：解码器的最后一层，将解码器输出的特征向量映射到目标词汇表的维度，通过 softmax 函数计算每个词的生成概率，从而得到最终的输出结果，如翻译后的句子。
• 信息传递
  • 编码器对源序列进行编码，将其转换为包含丰富语义和位置信息的特征表示，然后通过信息传递将这些特征传递给解码器。解码器利用编码器的输出，结合自身的掩码多头注意力和多头注意力机制，逐步生成目标序列。

多头注意力

• 对同一key，value，query，希望抽取不同的信息
  • 例如短距离关系和长距离关系
• 多头注意力使用h个独立的注意力池化
  • 合并各个头（head）输出得到最终输出
    

有掩码的多头注意力

• 解码器对序列中一个元素输出时，不应该考虑该元素之后的元素
• 可以通过掩码来实现
  • 生成掩码矩阵：通常会根据任务需求生成一个掩码矩阵，该矩阵的维度与输入序列的维度相匹配。在矩阵中，需要屏蔽的位置被设置为一个特定的值（如负无穷），而允许关注的位置则设置为 0 或其他合适的值。
  • 与注意力计算结合：在计算多头注意力的过程中，将掩码矩阵与注意力得分矩阵进行相应的运算（如相加）。这样，在进行 softmax 操作计算注意力权重时，被掩码的位置的得分会趋近于 0，从而在加权求和时不会对结果产生影响。

基于位置的前馈网络

• 基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的（positionwise）的原因。
• 将输入形状由（b, n, d）变换为（bn，d）
• 作用两个全连接层(1*1卷积层)
• 输出形状由（bn，d）变化回（b, n, d）
• 等价于两层核窗口为1的一维卷积层

层归一化

• 批量归一化对每个特征/通道里元素进行归一化
  • 不适合序列长度会变得NLP应用
• 层归一化对每个样本里的元素进行归一化；层归一化通过对每一层的输入数据进行归一化，将其均值和方差调整到固定值，使得数据分布更加稳定，从而可以使用更大的学习率，加快模型的收敛速度，减少训练时间。

信息传递

• 假设：编码器中的输出y1,…, y_n
• 将其作为解码中第i个Transformer块中多头注意力的key和value
  • 它的query来自目标序列
• 意味着编码器和解码器中块的个数和输出维度都是一样的。

在编码器完成对输入序列的编码后，编码器的输出会作为解码器中交叉注意力机制的输入之一。解码器通过交叉注意力机制计算与编码器输出的注意力权重，有选择地从编码器输出中提取相关信息，将输入序列的信息传递到解码器中，从而指导解码器生成符合输入语义的输出序列。

总结

• Transformer是一个纯使用注意力的编码-解码器
• 编码器和解码器都有n个transformer块
• 每个块里使用多头（自）注意力，基于位置的前馈网络，和层归一化。

代码实现

• 多头注意力机制实现

# 多头注意力

import math
import torch
from torch import nn
from d2l import torch as d2l

class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状:
        # (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状:
        # (batch_size，)或(batch_size，查询的个数)
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数，
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

# 使多个头并行计算， 上面的MultiHeadAttention类将使用下面定义的两个转置函数。 具体来说，transpose_output函数反转了transpose_qkv函数的操作。
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

# 测试
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape

• transformer实现

# transformer

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

# 基于位置的前馈网络
# 基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的（positionwise）的原因
# 输入X的形状：（批量大小，时间步数或序列长度，隐单元数或特征维度）
# 被一个两层的感知机转换为形状为（批量大小，时间步数，ffn_num_outputs）
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))), ffn(torch.ones((2, 3, 4)))[0]

# 实现残差连接和层规范化
# 层规范化和批量规范化的目标相同，但层规范化是基于特征维度进行规范化。
# 尽管批量规范化在计算机视觉中被广泛应用 
# 但在自然语言处理任务中（输入通常是变长序列）批量规范化通常不如层规范化的效果好。
# 对比不同维度的层规范化和批量规范化的效果
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))

# [使用残差连接和层规范化]来实现AddNorm
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

# 编码器
# 实现编码器的一个层
# 包含两个子层：多头自注意力和基于位置的前馈网络，这两个子层都使用了残差连接和紧随的层规范化。
"""
参数:
    key_size (int): 键（key）向量的维度大小。
    query_size (int): 查询（query）向量的维度大小。
    value_size (int): 值（value）向量的维度大小。
    num_hiddens (int): 隐藏层的维度大小，用于多头注意力机制和前馈神经网络。
    norm_shape (tuple): 层归一化（Layer Normalization）操作的形状。
    ffn_num_input (int): 前馈神经网络的输入维度。
    ffn_num_hiddens (int): 前馈神经网络隐藏层的维度。
    num_heads (int): 多头注意力机制中的头数。
    dropout (float): 丢弃率，用于防止过拟合。
    **kwargs: 其他关键字参数，传递给父类的构造函数。
"""
class EncoderBlock(nn.Module):
    """编码器块"""
    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

# 实现编码器的代码： 通过堆叠num_layers个EncoderBlock
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)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X

# 下面，指定超参数，创建一个两层的Transformer编码器
# 编码器输出形状：批量大小，时间步数目，num_hiddens）
encoder = TransformerEncoder(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

# 解码器
# transformer 解码器也是由多个相同的层组成
# 包含三个子层：解码器自注意力机制、“编码器-解码器”注意力 和 基于位置的前馈网络
# 子层也都被残差连接和紧随的层规范化围绕
class DecoderBlock(nn.Module):
    """
    解码器中第 i 个块的实现; i (int): 当前解码器块的索引。
    """
    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):
        """
        定义解码器块的前向传播逻辑。
        参数:
            X (torch.Tensor): 输入张量，形状为 (batch_size, num_steps, num_hiddens)。
            state (tuple): 包含编码器输出和相关信息的元组，具体格式为 (enc_outputs, enc_valid_lens, [None] * num_blocks)。
        返回:
            tuple: 包含解码器块输出和更新后的状态的元组。
        """
        # 从 state 中提取编码器的输出和有效长度
        enc_outputs, enc_valid_lens = state[0], state[1]
        # 训练阶段，输出序列的所有词元都在同一时间处理，因此 state[2][self.i] 初始化为 None。
        # 预测阶段，输出序列是通过词元一个接着一个解码的，因此 state[2][self.i] 包含着直到当前时间步第 i 个块解码的输出表示
        if state[2][self.i] is None:
            # 如果 state[2][self.i] 为 None，说明是训练阶段或者刚开始预测，将输入 X 作为键值对
            key_values = X
        else:
            # 否则，将之前解码的输出和当前输入 X 拼接起来作为键值对
            key_values = torch.cat((state[2][self.i], X), axis=1)
        # 更新 state[2][self.i] 为当前的键值对
        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)
        # 通过残差连接和层归一化更新输入 X
        Y = self.addnorm1(X, X2)
        # 编码器-解码器注意力机制
        # 计算编码器-解码器注意力的输出
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        # 通过残差连接和层归一化更新输入 Y
        Z = self.addnorm2(Y, Y2)
        # 前馈神经网络
        # 将 Z 输入到前馈神经网络中进行非线性变换，然后通过残差连接和层归一化得到最终输出
        return self.addnorm3(Z, self.ffn(Z)), state

# 编码器－解码器”注意力中进行缩放点积计算和残差连接中进行加法计算，[编码器和解码器的特征维度都是num_hiddens
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


# Transformer解码器：由num_layers个DecoderBlock组成
class TransformerDecoder(d2l.AttentionDecoder):
    """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, **kwargs):
        """初始化 Transformer 解码器的各个组件。"""
        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):
        """
        初始化解码器的状态
        参数:
            enc_outputs (torch.Tensor): 编码器的输出，形状为 (batch_size, num_steps, num_hiddens)。
            enc_valid_lens (torch.Tensor): 编码器输出的有效长度，用于掩码操作。
        返回:
            list: 包含编码器输出、有效长度和初始化为 None 的解码器块状态列表的列表。
        """
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]

    def forward(self, X, state):
        """
        定义解码器的前向传播逻辑
        参数:
            X (torch.Tensor): 输入的词元索引张量，形状为 (batch_size, num_steps)。
            state (list): 解码器的状态，由 init_state 方法初始化。
        返回:
            tuple: 包含解码器的最终输出和更新后的状态的元组。
        """
        # 将输入的词元索引转换为嵌入向量，并乘以隐藏层维度的平方根进行缩放
        # 然后添加位置编码信息
        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):
        """
        获取解码器的注意力权重。
        返回:
            list: 包含解码器自注意力和编码器 - 解码器注意力权重的列表。
        """
        return self._attention_weights

# 训练
# 指定Transformer的编码器和解码器都是2层，都使用4头注意力
# 下面使用“英语－法语”机器翻译数据集上训练Transformer模型

num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 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)

训练结果：

# 训练结束，使用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}')

结果：

# 可视化权重
enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
    -1, num_steps))
enc_attention_weights.shape

结果：

# 可视化多头注意力机制中不同头（head）的注意力权重矩阵
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))

# 可视化解码器的自注意力权重和“编码器－解码器”的注意力权重
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
# 由于解码器自注意力的自回归属性，查询不会对当前位置之后的“键－值”对进行注意力计算。
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))