Pytorch attention transformer整理

一、函数运用

1.torch.eye(n), n为维数

生成n*n的对角矩阵

2.torch.rand(n),n为数量

产生n个从[0,1)内均匀分布的数值

3.torch.sort(xxx) 给产生的张量进行排序

该函数返回两个值，但我们只关心第一个，所以

x_train, _ = torch.sort(torch.rand(n_train) * 5) 

4.torch.normal 正态分布

生成一个均值为0，方差为0.5的，并且大小为n_train的张量，注意，定义大小默认的参数类型是元组，(n_train,)表示单个元素的元组

torch.normal(0.0, 0.5, (n_train,))

5.torch.arrange(a,b,k) 创建张量

创建一个a到b，步长为k的张量
类比range

x_test = torch.arange(0, 5, 0.1)

6.torch.interleave 重复元素

repeat_interleave(n_train): 这个函数将 x_test 中的每个元素重复 n_train 次。例如，如果 x_test 包含 [0.1, 0.2, 0.3]，而 n_train 为 3，则结果将是 [0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.3, 0.3, 0.3]。
reshape((-1, n_train))-1代表让函数自动推断，形成一个（？，n_train）的张量

X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))

7.torch.matmul 矩阵乘法

y_hat = torch.matmul(attention_weights, y_train)

torch.matmul(attention_weights, y_train) 将对这两个张量进行矩阵乘法操作。具体来说，如果 attention_weights 是形状为 (a, b) 的矩阵，y_train 是形状为 (b, c) 的矩阵，那么乘积的结果将是一个形状为 (a, c) 的矩阵，其中的每个元素是矩阵乘积的结果。

8.torch.bmm 批量矩阵乘法

例如n个Xi矩阵，维数是1行4列，n个Yi矩阵，维数是4行6列，相乘为n个1行6列的张量

X = torch.ones((2, 1, 4))
Y = torch.ones((2, 4, 6))
torch.bmm(X, Y).shape

9.torch.ones(z,x,y)

产生z个x行y列的张量

X = torch.ones((2, 1, 4))
Y = torch.ones((2, 4, 6))

10.tensor(张量即可).unsqueeze() 添加一维维度

weights.unsqueeze(1) 是一个 PyTorch 或者类似的张量操作。这个操作是在张量中添加一个维度，具体来说是在索引为 1 的位置（从 0 开始计数）上添加一个新的维度。这个新维度的大小是 1。

假设 weights 是一个形状为 (n,) 的张量，即一维张量。那么 weights.unsqueeze(1) 将会把这个一维张量变成一个形状为 (n, 1) 的二维张量。新的维度在原张量的索引 1 的位置上添加，其他维度的大小保持不变。

weights = torch.ones((2, 10)) * 0.1
values = torch.arange(20.0).reshape((2, 10))
torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))

weight张量从原来的（2，10）变为（2，1，10）
values张量从原来的（2，10）变为（2，10，1）
产生的结果为（2，1，1）

11.tensor(张量即可).repeat()

x_train.repeat((n_train, 1))

在第一个维度，重复n_train次

12.torch.repeat_interleave(valid_lens, shape[1])

具体来说，torch.repeat_interleave(valid_lens, shape[1]) 的作用是将 valid_lens 中的每个整数值重复插入到输出张量中，重复次数为 shape[1] 次。这样，输出张量的长度将会是 valid_lens 的总和乘以 shape[1]。

例如，如果 valid_lens 是一个长度为 batch_size 的一维张量，而 shape[1] 是一个常数 n，那么输出张量的长度将会是 batch_size * n。每个 valid_lens 中的整数值都会在输出张量中重复插入 n 次。

valid_lens = torch.repeat_interleave(valid_lens, shape[1])

13.net.eval() 评估模式

attention.eval() 是将 AdditiveAttention 模型设置为评估模式。在 PyTorch 中，模型在训练时和评估时有不同的行为，例如 Dropout 和 Batch Normalization 在训练时会起作用，而在评估时会关闭。 通过调用 .eval() 方法，我们将模型设置为评估模式，这样在进行推理时就不会产生随机性。

attention.eval()

14.tensor（张量）.transpose()交换维度

当我们调用 keys.transpose(1, 2) 时，它将交换第一个维度和第二个维度的位置。 具体地说，它将原来的形状 (batch_size, num_kv_pairs, num_hidden) 转换为 (batch_size, num_hidden, num_kv_pairs)。这样做的效果是将键值对的特征维度移动到了第二个维度的位置，使得每个键值对的特征维度可以在进行矩阵乘法时与其他张量的维度相匹配。

keys.transpose(1,2)

15.tensor（张量）.permute()

X.permute(0, 2, 1, 3) 是 PyTorch 中的张量操作，用于对张量的维度进行重新排列。具体来说，它将张量的维度按照指定的顺序进行排列。

X = X.permute(0, 2, 1, 3)

二、 nn模块

1.softmax

nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)

dim=1，规定为在行上进行运算
dim=-1，在最后一个维度进行计算

2.nn.MSELoss（） 计算均方误差

loss = nn.MSELoss(reduction='none')

参数reduction用于指定损失的计算方式，有三种选项：
“none”：表示不进行任何降维，即每个样本的损失都会被保留，不进行汇总。这种情况下，返回的损失张量的形状与输入张量的形状相同，每个样本对应一个损失值。
“mean”：表示对每个样本的损失进行求均值，返回的损失值是所有样本损失的平均值。
“sum”：表示对每个样本的损失进行求和，返回的损失值是所有样本损失的总和。

3.nn.optim.SGD()随机梯度下降法

运用随机梯度下降优化器来更新神经网络的参数。SGD优化器在每次迭代时只使用一个样本（或一批样本）来计算梯度，并根据该梯度更新参数。

trainer = torch.optim.SGD(net.parameters(), lr=0.5)

指定优化参数，net.parameters()，定义学习率（步长），学习率是一个控制参数更新步长的超参数。它决定了每次参数更新的大小，对应着优化器中梯度乘以的比例。较大的学习率意味着更大的参数更新步长，而较小的学习率则意味着更小的步长。在这里，学习率被设置为0.5，这是一个相对较大的值，但实际上最优的学习率取决于具体的任务和数据。

4.nn.Linear()定义线性变化模块

nn.Linear 接受两个参数：

in_features：输入特征的数量，即输入张量的大小（除了批次维度之外的所有维度的乘积）。
out_features：输出特征的数量，即输出张量的大小（除了批次维度之外的所有维度的乘积）。
还可以通过bias=False 表示不使用偏置参数，即不向线性变换中添加偏置。

nn.Linear(key_size, num_hiddens, bias=False)

5.nn.Dropout()正则化技术，用于防止神经网络过拟合

nn.Dropout(dropout) 的作用是在每次前向传播时，以概率 dropout 随机丢弃输入张量的某些元素，并对保留下来的元素进行缩放，以保持期望的输出值的期望值不变。这样可以模拟出多个不同的子网络，从而提高模型的鲁棒性。

nn.Dropout(dropout)

6.nn.Parameter 初始化网络参数

self.w = nn.Parameter(torch.rand((1,), requires_grad=True))

torch.rand((1,))产生形状为（1，）的张量，并需要计算梯度

7.层规范化和批量规范化

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))

层规范化主要基于特征维度进行规范化，批量规范化主要基于每个batch进行规范化
即层规范化就是一层一层进行规范化，（每个元素-平均值）/标准差
批量规范化则为按照batch维度，即不同的batch进行规范化，可以理解为按照列来
在自然语言处理任务中（输入通常是变长序列）批量规范化通常不如层规范化的效果好
注意：
1.在 PyTorch 中，nn.LayerNorm 的输入张量的最后一个维度确定了沿着哪个维度进行归一化

例如，输入张量的形状为 (N, D)：

在这种情况下，LayerNorm 将对每个样本的特征维度进行归一化。 例如，如果 N=2，表示有两个样本，D=3，表示每个样本有三个特征，那么 LayerNorm 将对每个样本的三个特征维度进行归一化。
2. nn.BatchNorm1d，它用于一维数据，通常用于处理批次数据。它会沿着第一个维度（即批次维度） 计算均值和方差，并对整个批次的数据进行归一化。

8.nn.Sequential()

nn.Sequential() 是一个容器，用于按顺序组合多个层或模块。

self.blks = nn.Sequential(
            nn.Linear(10, 20),
            nn.ReLU(),
            nn.Linear(20, 10),
            nn.ReLU()
        )

9 add_module() 向模型中添加子模块

用于向模型中添加子模块（submodules）。它允许在现有模块的容器中添加子模块，并为这些子模块指定一个名称。使用 add_module 方法，你可以向 nn.Sequential() 容器中添加子模块，并为每个子模块指定一个名称。这样做的好处是可以在模型中对每个子模块进行命名，便于后续对子模块进行调用或查找。

container.add_module(name, module)

10.nn.Embedding()创建嵌入层

num_embeddings：表示词汇表的大小，即总共有多少个不同的单词或类别。
embedding_dim：表示嵌入向量的维度，即每个单词或类别被映射为一个多少维的向量。

import torch
import torch.nn as nn

# 创建一个词汇表大小为10，词嵌入维度为5的嵌入层
embedding = nn.Embedding(num_embeddings=10, embedding_dim=5)

# 输入一个大小为3的批次，每个样本包含两个单词的序列
input_tensor = torch.LongTensor([[1, 2], [3, 4], [5, 6]])

# 将输入数据传递给嵌入层进行映射
output = embedding(input_tensor)

print(output)
# 输出的张量大小为(3, 2, 5)，表示3个样本，每个样本包含2个单词，每个单词映射为一个5维的向量

11.num_steps指每个序列的时间步数

num_steps 参数定义了一个序列的最大长度。如果输入序列或输出序列的长度超过了这个值，那么就会进行截断或者填充（padding）操作，以保证所有序列具有相同的长度。这种处理通常用于简化模型的实现和提高训练效率。
举个例子，如果 num_steps 设置为 10，那么任何超过 10 个单词的源语言句子或目标语言句子都会被截断或填充到长度为 10 的序列。这样可以确保每个批次中的所有序列具有相同的长度，方便模型的并行计算和训练。

定义网络

（1）定义类，继承nn.module
（2）初始化网络参数 nn.Parameter
（3）定义前向传播网络
（4）定义数据，例如注意力机制需要将keys，values的对角线元素去除，可以确保在计算注意力权重时不考虑每个元素与自身的相关性，从而避免将自身作为参考对象，使注意力权重更加准确地反映了元素之间的关联程度。
（5）构建网络
（6）定义损失函数
（7）定义优化函数
（8）模型训练

例：

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])

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() #l包含每个样本的损失值，求和并进行反向传播更新梯度
    trainer.step()  # 使用优化器更新模型参数
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

trainer.step()
调用优化器的 step() 方法来更新模型的参数。优化器将根据梯度信息和学习率来更新模型的参数。

三、注意力机制（整理自李沐的b站课程）

1.注意力汇聚

上述为注意力汇聚公式（attention pooling）。x是查询，（xi，yi）是键值对，注意力汇聚就是yi的加权平均。 将查询和键之间的关系建模为 注意力权重（attention weight） a(x,xi)。
例如，高斯核

键xi越接近给定的x，那么分配给这个键对应的yi的注意力权重就会很大

2.attention机制原理

高斯核指数部分可以视为注意力评分函数（attention scoring function）， 简称评分函数（scoring function）， 然后把这个函数的输出结果输入到softmax函数中进行运算。 通过上述步骤，将得到与键对应的值的概率分布（即注意力权重）。 最后，注意力汇聚的输出就是基于这些注意力权重的值的加权和。

3.注意力评分函数

（1）掩码softmax

为了仅将有意义的词元作为值来获取注意力汇聚， 可以指定一个有效序列长度（即词元的个数）， 以便在计算softmax时过滤掉超出指定范围的位置，其中任何超出有效长度的位置都被掩蔽并置为0。

def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1) #相当于在列上进行softmax操作
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6) #采用d2l库中的函数，将超出范围的置换成很大的数，从而softmax为0
        return nn.functional.softmax(X.reshape(shape), dim=-1) #在最后一个维度上进行softmax函数计算

解释一下，reshape（-1，shape[-1]）
在PyTorch的reshape函数中，参数-1表示根据其他维度的大小和张量的总元素数量来自动推断该维度的大小。 这种用法非常方便，尤其是在需要根据张量的总元素数量来计算另一个维度的大小时。

例如，假设有一个张量x，它的形状是(2, 3, 4)，即有2个维度为3x4的矩阵。如果我们执行x.reshape(-1, 4)，PyTorch会自动计算第一个维度的大小，使得张量的总元素数量保持不变，即2x3x4=24。因此，第一个维度的大小将被设置为24/4=6，结果张量的形状将变为(6, 4)。

（2）加性注意力，适用于查询和键是不同长度的矢量

class AdditiveAttention(nn.Module):
    """加性注意力"""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在维度扩展后，
        # queries的形状：(batch_size，查询的个数，1，num_hidden)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        features = torch.tanh(features)
        # self.w_v仅有一个输出，因此从形状中移除最后那个维度。
        # scores的形状：(batch_size，查询的个数，“键-值”对的个数)
        scores = self.w_v(features).squeeze(-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # values的形状：(batch_size，“键－值”对的个数，值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)

解释一下features怎么计算得到的
features = queries.unsqueeze(2) + keys.unsqueeze(1)：在这里，queries 和 keys 的形状都被扩展了两个维度，以便将它们进行广播相加。具体来说，queries.unsqueeze(2) 将 queries 张量的维度扩展为 (batch_size, 查询的个数, 1, num_hidden)，而 keys.unsqueeze(1) 将 keys 张量的维度扩展为 (batch_size, 1, “键－值”对的个数, num_hiddens)。然后，通过广播操作，两个张量的每个对应位置进行相加，得到了一个形状为 (batch_size, 查询的个数, “键－值”对的个数, num_hidden) 的张量 features。

（3）缩放点积注意力机制，当键和查询的长度相同时计算效率更高

class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 设置transpose_b=True为了交换keys的最后两个维度
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)

4.多头注意力机制

与其只使用单独一个注意力汇聚， 我们可以用独立学习得到的h组不同的 线性投影（linear projections）来变换查询、键和值。 然后，这h组变换后的查询、键和值将并行地送到注意力汇聚中。 最后，将这h个注意力汇聚的输出拼接在一起， 并且通过另一个可以学习的线性投影进行变换， 以产生最终输出。 这种设计被称为多头注意力。

采用的缩放点积注意力
注：为了避免计算代价和参数代价的大幅增长，设定$p_q=p_k=p_v=p_o/h$。值得注意的是，如果将查询、键和值的线性变换的输出数量设置为$p_{q}h=p_{k}h=p_{v}h=p_o$，则可以并行计算$h$个头。在下面的实现中，$p_o$是通过参数num_hiddens指定的。

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)
   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])

	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)

最终输出的维度是(batch_size，查询的个数，num_hiddens)

5.自注意力和位置编码

(1)自注意力

将词元序列输入注意力池化中， 以便同一组词元同时充当查询、键和值。 具体来说，每个查询都会关注所有的键－值对并生成一个注意力输出。 由于查询、键和值来自同一组输入，因此被称为 自注意力（self-attention）。

基于多头注意力对一个张量完成自注意力的计算， 张量的形状为（批量大小，时间步的数目或词元序列的长度d）。 输出与输入的张量形状相同。

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, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
attention(X, X, X, valid_lens).shape

（2）位置编码

class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)

self.P[:, :X.shape[1], :]：这一步是从位置编码矩阵 self.P 中选取与输入序列长度相匹配的位置编码。具体来说，:X.shape[1] 是取输入张量 X 的序列长度，即 seq_length，然后 self.P[:, :X.shape[1], :] 就是从位置编码矩阵中选取前 seq_length 个位置的编码。这样做是因为输入张量的长度可能小于位置编码矩阵的长度，所以需要进行截断。

四、Transformer

1.基于位置的前馈网络

输入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)))

2.残差连接和层规范化

残差连接的基本思想是在网络的某些层之间添加一个跨层的直接连接。这个直接连接可以绕过某些层，将输入信号直接传递给网络的后续层。这样，网络可以学习将输入信号变换为输出信号的增量，而不是直接学习将输入映射到输出。

具体来说，残差连接的公式如下所示：
Output = Activation( Input + F(Input))
其中，Input 表示输入信号，F(Input) 表示某一层的变换操作，Activation 表示激活函数。残差连接的关键在于直接将输入信号与变换后的输出信号相加，并将其传递给激活函数。这样，即使变换操作无法有效地学习到输入与输出之间的映射关系，也可以通过残差连接直接将输入信号传递给输出，从而保证信息的流动。

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)

3.编码块

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))

Transformer编码器中的任何层都不会改变其输入的形状

4.Transformer 编码器

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

5.解码块

在掩蔽多头解码器自注意力层（第一个子层）中，查询、键和值都来自上一个解码器层的输出。关于序列到序列模型（sequence-to-sequence model），在训练阶段，其输出序列的所有位置（时间步）的词元都是已知的；然而，在预测阶段，其输出序列的词元是逐个生成的。因此，在任何解码器时间步中，只有生成的词元才能用于解码器的自注意力计算中。为了在解码器中保留自回归的属性，其掩蔽自注意力设定了参数dec_valid_lens，以便任何查询都只会与解码器中所有已经生成词元的位置（即直到该查询位置为止）进行注意力计算。

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 = MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = 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]
        # 训练阶段，输出序列的所有词元都在同一时间处理，
        # 因此state[2][self.i]初始化为None。
        # 预测阶段，输出序列是通过词元一个接着一个解码的，
        # 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示
        if state[2][self.i] is None:
            key_values = X
        else:
            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

        #**自注意力  training的qkv一样，但是predict的kv是结合之前输出的qv**
        
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # 编码器－解码器注意力。
        # enc_outputs的开头:(batch_size,num_steps,num_hiddens)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens) #来自encoder的kv
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

6.transformer 解码器

最后，通过一个全连接层计算所有vocab_size个可能的输出词元的预测值。解码器的自注意力权重和编码器解码器注意力权重都被存储下来，方便日后可视化的需要。

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

注意：在原始的 Transformer 模型中，每个位置的位置编码值的范围在 −1 到 1之间，因此需要乘以一个缩放因子。常见的缩放因子是sqrt(d_model),其中d_model是词嵌入的维度。这个缩放因子有助于保持位置编码的数值范围与词嵌入的数值范围相匹配，使得它们在加法时不会导致数值过大或过小。