【深度学习】注意力机制（Transformer）

注意力机制

1.基础概念

1.1 查询、键和值

在人类的注意力方式中，有自主性的与非自主性的注意力提示两种解释方式。所谓自主性注意力提示，就是人本身主动想要关注到的某样东西；非自主性提示则是基于环境中物体的突出性和易见性，由客观因素影响使得人不自知地而关注到某样东西。在注意力机制中也是通过这两种注意力提示在神经网络中搭建相关框架。

首先，考虑一个相对简单的状况，即只用非自主性提示。要想选择偏向于感官输入，则可以简单地使用参数化的全连接层，甚至是非参数化的最大汇聚层或者平均汇聚层。因此，"是否包含自主性提示"将注意力机制与全连接层或汇聚层区别开来。 在注意力机制的背景下，自主性提示被称为查询（query） 。给定任何查询，注意力机制通过注意力汇聚（attention pooling）将选择引导至感官输入（sensory inputs）。在注意力机制中，这些感官输入被称为值（value） 。更通俗的解释，每个值都与一个键（key）配对，这可以想象为感官输入的非自主提示。如下图所示，可以通过设计注意力汇聚的方式，便于给定的查询（自主性提示）与键（非自主性提示）进行匹配，这将引导得出最匹配的值（感官输入） 。

1.2 注意力的可视化

平均汇聚层可以被视为输入的加权平均值，其中各输入的权重是一样的。实际上，注意力汇聚得到的是加权平均的总和值，其中权重是在给定的查询和不同的键之间计算得出的。

为了可视化注意力权重，需要定义一个show_heatmaps函数。其中输入matrices的形状是 （要显示的行数，要显示的列数，查询的数目，键的数目）。代码如下：

import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable

def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5), cmap='Reds'):
    """显示矩阵热图"""
    num_rows, num_cols = matrices.shape[0], matrices.shape[1]
    fig, axes = plt.subplots(num_rows, num_cols, figsize=figsize, sharex=True, sharey=True, squeeze=False)

    for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):
        for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):
            pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap)
            if i == num_rows - 1:
                ax.set_xlabel(xlabel)
            if j == 0:
                ax.set_ylabel(ylabel)
            if titles:
                ax.set_title(titles[j])
        fig.colorbar(pcm, ax=axes, shrink=0.6)

attention_weights = torch.eye(10).reshape((1, 1, 10, 10))
show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

输出得到：

2.注意力汇聚：Nadaraya-Watson核回归

注意力机制的主要框架为：查询（自主提示）和键（非自主提示）之间的交互形成了注意力汇聚；注意力汇聚有选择地聚合了值（感官输入）以生成最终的输出结果。1964年提出的Nadaraye-Watson核回归模型是一个简单且完整的体现注意力机制在实践中的运作的例子，可以演示具有注意力机制的机器学习。

2.1 生成数据集

简单起见，考虑下面这个回归问题：给定的成对的"输入-输出"数据集 ( x 1 , y 1 ) , ... , ( x n , y n ) {(x_1,y_1), \dots , (x_n,y_n)} (x1,y1),...,(xn,yn)， 如何学习 f f f 来预测任意新输入 x x x的输出 y ^ = f ( x ) \hat{y} = f(x) y^=f(x) ？

根据下面的非线性函数生成一个人工数据集，其中加入的噪声项为 ϵ \epsilon ϵ ：
y i = 2 s i n ( x i ) + x i 0.8 + ϵ y_i = 2sin(x_i) + x_i^{0.8} + \epsilon yi=2sin(xi)+xi0.8+ϵ

其中 ϵ \epsilon ϵ 服从均值为0和标准差为0.5的正态分布。这里生成50个训练样本和50个测试样本。为了更好地可视化之后的注意力模式，需要将训练样本进行排序。代码如下：

import torch
from torch import nn

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.5, (n_train,)) #训练样本的输出
x_test = torch.arange(0, 5, 0.1) #测试样本
y_truth = f(x_test) #测试样本的真实输出
n_test = len(x_test)
n_test

运行代码后可以得到输出为50.下面定义函数来绘制所有训练样本（样本由圆圈表示），不带噪声项的真实数据生成函数f（标记为"truth"），以及学习得到的预测函数（标记为"pred"）。

import matplotlib.pyplot as plt

#绘制训练样本
def plot_kernel_reg(x_test, y_truth, y_hat, x_train=None, y_train=None):
    # plt.figure(figsize=(8, 6))
    plt.plot(x_test, y_truth, '-', label='Truth')
    plt.plot(x_test, y_hat, '-', label='Pred')
    
    if x_train is not None and y_train is not None:
        plt.plot(x_train, y_train, 'o', alpha=0.5, label='Train data')
    
    plt.legend()
    plt.grid(True, linestyle='--', color='gray')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.show()

2.2 平均汇聚

先使用最简单的估计器来解决回归问题。基于平均汇聚来计算所有训练样本输出值的平均值：
f ( x ) = 1 n ∑ i = 1 n y i f(x) = \frac{1}{n} \sum_{i=1}^n y_i f(x)=n1i=1∑nyi

y_hat = torch.repeat_interleave(y_train.mean(), n_test)
plot_kernel_reg(x_test, y_truth, y_hat, x_train, y_train)

输出结果如下图所示，可以看到这个估计器还不够聪明，真实函数和预测函数相差很大。

2.3 非参数注意力汇聚

显然，平均汇聚忽略了输入 x i x_i xi。于是Nadaraya和Waston提出了一个更好的想法，即根据输入的位置对输出 y i y_i yi进行加权：
f ( x ) = ∑ i = 1 n K ( x − x i ) ∑ j = 1 n K ( x − x j ) y i f(x)=\sum_{i=1}^{n} \frac{K\left(x-x_{i}\right)}{\sum_{j=1}^{n} K\left(x-x_{j}\right)} y_{i} f(x)=i=1∑n∑j=1nK(x−xj)K(x−xi)yi

其中K是核（kernal）。公式所描述的估计器被称为Nadaraya-Waston核回归。由此我们可以在注意力机制框架的角度重写一个更加通用的注意力汇聚公式：
f ( x ) = ∑ i = 1 n α ( x , x i ) y i f(x) = \sum_{i=1}^n \alpha (x,x_i)y_i f(x)=i=1∑nα(x,xi)yi

其中x是查询， ( x i , y i ) (x_i, y_i) (xi,yi)是键值对。此时注意力汇聚是 y i y_i yi的加权平均。将查询x和 x i x_i xi之间的关系建模为注意力权重（attention weight） α ( x , x i ) \alpha (x,x_i) α(x,xi)，这个权重将被分配给每一个对应值 y i y_i yi。对于任何查询，模型在所有键值对注意力权重都是一个有效的概率分布：它们是非负的，并且总和为1。

为了更好地理解注意力汇聚，下面考虑一个高斯核（Gauessian kernal），其定义为：
K ( u ) = 1 2 π e x p ( − u 2 2 ) K(u) = \frac{1}{\sqrt{2 \pi} } exp(- \frac{u^2}{2}) K(u)=2π 1exp(−2u2)

将高斯核代入可以得到：
f ( x ) = ∑ i = 1 n α ( x , x i ) y i = ∑ i = 1 n exp ⁡ ( − 1 2 ( x − x i ) 2 ) ∑ j = 1 n exp ⁡ ( − 1 2 ( x − x j ) 2 ) y i = ∑ i = 1 n softmax ⁡ ( − 1 2 ( x − x i ) 2 ) y i . \begin{aligned} f(x) & =\sum_{i=1}^{n} \alpha\left(x, x_{i}\right) y_{i} \\ & =\sum_{i=1}^{n} \frac{\exp \left(-\frac{1}{2}\left(x-x_{i}\right)^{2}\right)}{\sum_{j=1}^{n} \exp \left(-\frac{1}{2}\left(x-x_{j}\right)^{2}\right)} y_{i} \\ & =\sum_{i=1}^{n} \operatorname{softmax}\left(-\frac{1}{2}\left(x-x_{i}\right)^{2}\right) y_{i} . \end{aligned} f(x)=i=1∑nα(x,xi)yi=i=1∑n∑j=1nexp(−21(x−xj)2)exp(−21(x−xi)2)yi=i=1∑nsoftmax(−21(x−xi)2)yi.

在上式中，如果一个键 x i x_i xi越是接近给定的查询x，那么分配给这个键对应值 y i y_i yi的注意力权重也就越大，这样就获得了更多的注意力。

值得注意的是，Nadaraya-Waston核回归是一个非参数模型，因此上式是非参数的注意力汇聚模型。接下来我们将基于这个非参数的注意力汇聚模型来绘制预测结果，从绘制的结果会发现新的模型预测线是平滑的，并且比平均汇聚的预测更接近真实。

#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(x_test, y_truth, y_hat, x_train, y_train)

现在来观察注意力的权重，这里测试数据的输入相当于查询，而训练数据的输入相当于键。因为两个输入都是排序过的，因此由观察可知"查询-键"对越接近，注意力汇聚的注意力权重就越高。

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

2.4 带参数注意力汇聚

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

例如，与高斯核注意力汇聚不同，在下面的查询和键之间的距离乘以可学习参数 ω \omega ω：
f ( x ) = ∑ i = 1 n α ( x , x i ) y i = ∑ i = 1 n exp ⁡ ( − 1 2 ( ( x − x i ) w ) 2 ) ∑ j = 1 n exp ⁡ ( − 1 2 ( ( x − x j ) w ) 2 ) y i = ∑ i = 1 n softmax ⁡ ( − 1 2 ( ( x − x i ) w ) 2 ) y i . \begin{aligned} f(x) & =\sum_{i=1}^{n} \alpha\left(x, x_{i}\right) y_{i} \\ & =\sum_{i=1}^{n} \frac{\exp \left(-\frac{1}{2}\left(\left(x-x_{i}\right) w\right)^{2}\right)}{\sum_{j=1}^{n} \exp \left(-\frac{1}{2}\left(\left(x-x_{j}\right) w\right)^{2}\right)} y_{i} \\ & =\sum_{i=1}^{n} \operatorname{softmax}\left(-\frac{1}{2}\left(\left(x-x_{i}\right) w\right)^{2}\right) y_{i} . \end{aligned} f(x)=i=1∑nα(x,xi)yi=i=1∑n∑j=1nexp(−21((x−xj)w)2)exp(−21((x−xi)w)2)yi=i=1∑nsoftmax(−21((x−xi)w)2)yi.

接下来将通过训练上述模型来学习注意力汇聚的参数。

2.4.1 批量矩阵乘法

为了更有效地计算小批量数据的注意力，我们可以利用深度学习开发框架中提供的批量矩阵乘法。

假设第一个小批量数据包含n个矩阵 X 1 , ... X n \mathrm {X_1}, \dots \mathrm{X_n} X1,...Xn，形状为 a × b a \times b a×b， 第二个小批量数据包含n个矩阵 Y 1 , ... Y n \mathrm {Y_1}, \dots \mathrm{Y_n} Y1,...Yn，形状为 b × c b \times c b×c。它们的批量矩阵乘法得到n个矩阵 X 1 Y 1 , ... X n Y n \mathrm {X_1Y_1}, \dots \mathrm{X_nY_n} X1Y1,...XnYn，形状为 a × c a \times c a×c。因此，假定两个张量的形状分别是 ( n , a , b ) (n,a,b) (n,a,b)和 ( n , b , c ) (n,b,c) (n,b,c)，它们的批量矩阵乘法输出的形状为 ( n , a , c ) (n,a,c) (n,a,c)。

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

输出：

torch.Size([2, 1, 6])

在注意力机制的背景中，我们可以使用小批量矩阵乘法来计算小批量数据中的加权平均值。

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

输出为：

tensor([[[ 4.5000]],

        [[14.5000]]])

2.4.2 定义模型

基于上一节的带参数的注意力汇聚，使用小批量矩阵乘法，定义Nadaraya-Watson核回归的带参数版本为：

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)

2.4.3 训练

接下来，将训练数据集变换为键和值用于训练注意力模型。在带参数的注意力汇聚模型中，任何一个训练样本的输入都会和除自己以外的所有训练样本的键值对进行计算，从而得到其对应的预测输出。

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

训练带参数的注意力汇聚模型时，使用平方损失函数和随机梯度下降。

from IPython import display
class Animator:
    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None, ylim=None, xscale='linear', yscale='linear',
                fmts=('-', 'm--', 'g-', 'r:'), nrows=1, ncols=1, figsize=(3.5, 2.5)):
        if legend is None:
            legend = []
            
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes,]

        # self.config_axes = lambda:self.set_axes(
        #     self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)
        self.X, self.Y, self.fmts = None, None, fmts
        self.set_axes(xlabel, ylabel, xlim, ylim, xscale, yscale, legend)

    def set_axes(self, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):
        """设置轴"""
        for ax in self.axes:
            ax.set_xlabel(xlabel)
            ax.set_ylabel(ylabel)
            ax.set_xscale(xscale)
            ax.set_yscale(yscale)
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)
            if legend:
                ax.legend(legend)
            ax.grid()
            
    def add(self, x, y):
        if not hasattr(y, '__len__'):
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)

        for ax in self.axes:
            ax.cla()
            for x, y, fmt in zip(self.X, self.Y, self.fmts):
                ax.plot(x, y, fmt)
        display.display(self.fig)
        display.clear_output(wait=True)

net = NWKernelRegression()
loss = nn.MSELoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = 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(x_test, y_truth, y_hat, x_train, y_train)

新的模型更不平滑了，下面看一下输出结果的绘制图：与非参数的注意力汇聚模型相比，带参数的模型加入可学习的参数后，曲线在注意力权重较大的区域变得更不平滑。

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

3.注意力评分函数

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

下图说明了如何将注意力汇聚的输出计算成为值的加权和，其中a表示注意力评分函数。由于注意力权重是概率分布，因此加权和其本质上是加权平均值。

3.1 掩蔽softmax操作

softmax操作用于输出一个概率分布作为注意力权重。在某些情况下，并非所有的值都应该被纳入到注意力汇聚中。例如在机器翻译中高效处理小批量数据集，某些文本序列被填充了没有意义的特殊词元，为了仅将有意义的词元作为值来获取注意力汇聚，可以指定一个有效序列长度（即词元的个数），以便在计算softmax时过滤掉超出指定范围的位置。下面的masked_softmax函数实现了这样的掩蔽softmax操作，其中任何超出有效长度的位置都被掩蔽并置为0.

def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X

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)
    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 = sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

考虑由两个 2 × 4 2 \times 4 2×4矩阵表示的样本，这两格样本的有效长度分别为2和3.经过掩蔽softmax操作，超出有效长度的值都被掩蔽为0.

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))

输出为：

tensor([[[0.5565, 0.4435, 0.0000, 0.0000],
         [0.5358, 0.4642, 0.0000, 0.0000]],

        [[0.4991, 0.2465, 0.2544, 0.0000],
         [0.3961, 0.3220, 0.2819, 0.0000]]])

3.2 加性注意力

一般来说，当查询和键是不同长度的矢量时，可以使用加性注意力作为评分函数。其原理是将查询和键连结起来后输入到一个多层感知机中，感知机包含一个隐藏层，其隐藏单元数是一个超参数h。通过tanh作为激活函数，并且禁用偏置项。

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)

用一个小例子来演示：

queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))
# values的小批量，两个值矩阵是相同的
values = torch.arange(40, dtype=torch.float32).reshape(1, 10, 4).repeat(
    2, 1, 1)
valid_lens = torch.tensor([2, 6])

attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8,
                              dropout=0.1)
attention.eval()
attention(queries, keys, values, valid_lens)

输出为：

tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)

尽管加性注意力包含了可学习的参数，但由于本例子中每个键都是相同的，所以注意力权重是均匀的，由指定的有效长度决定。

show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

3.3 缩放点积注意力

使用点积可以得到计算效率更高的评分函数，但是点积操作要求查询和键具有相同的长度d。假设查询和键的所有元素都是独立的随机变量，并且都满足零均值和单位方差，那么两个向量的点积的均值为0，方差为d。为确保无论向量长度如何，点积的方差在不考虑向量长度的情况下仍然是1，我们再将点积除以 d \sqrt{d} d ，则缩放点积注意力评分函数为：
a ( q , k ) = q ⊤ k / d a(\mathbf{q}, \mathbf{k})=\mathbf{q}^{\top} \mathbf{k} / \sqrt{d} a(q,k)=q⊤k/d

下面的缩放点积注意力的实现使用了暂退法进行模型正则化。

import math

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)

使用先前相同的例子进行操作，在点积操作中令查询的特征维度与键的特征维度大小相同。

queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
attention(queries, keys, values, valid_lens)

输出为：

tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]])

与加性注意力演示相同，由于键包含的是相同的元素， 而这些元素无法通过任何查询进行区分，因此获得了均匀的注意力权重。

show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

4.Bahdanau注意力

循环神经网络编码器可以将长度可变的序列转换为固定形状的上下文变量，然后循环神经网络解码器根据生成的词元和上下文变量按词元生成输出（目标）序列词元。然而，即使并非所有输入（源）词元都对解码某个词元都有用，在每个解码步骤中仍然使用编码相同的上下文变量。为了考虑如何改变上下文变量，在为给定文本序列生成手写的挑战中，Graves设计了一种可微注意力模型，将文本字符与更长的笔迹对齐，其中对齐方式仅向一个方向移动。受学习对齐想法的启发，Bahdanau等人提出了一个没有严格单向对齐限制的可微注意力模型。在预测词元时，如果不是所有输入词元都相关，模型将仅对齐（或参与）输入序列中与当前预测相关的部分。这是通过将上下文变量视为注意力集中的输出来实现的。

4.1 模型

假设输入序列有T个词元，解码时间步 t ′ t' t′的上下文变量是注意力集中的输出：
c t ′ = ∑ t = 1 T α ( s t ′ − 1 , h t ) h t \mathbf{c}{t'}= \sum{t=1}^T \alpha (\mathbf{s}_{t'-1}, \mathbf{h}_t) \mathbf{h}_t ct′=t=1∑Tα(st′−1,ht)ht

其中，时间步 t ′ − 1 t'-1 t′−1时的解码器隐状态 s t ′ − 1 \mathbf{s}_{t'-1} st′−1是查询，编码器隐状态 h t \mathbf{h}_t ht既是键，也是值，注意力权重 α \alpha α 是使用所定义的加性注意力打分函数计算的。

下图描述了Bahdanau注意力的架构。

4.2 定义注意力解码器

为了更方便地显示学习的注意力权重， 以下AttentionDecoder类定义了带有注意力机制解码器的基本接口。

import torch
from torch import nn

#解码器
class Decoder(nn.Module):
    """编码器-解码器架构的基本解码器接口"""
    def __init__(self, **kwargs):
        super(Decoder, self).__init__(**kwargs)

    def init_state(self, enc_outputs, *args):
        raise NotImplementedError

    def forward(self, X, state):
        raise NotImplementedError

class AttentionDecoder(Decoder):
    """带有注意力机制解码器的基本接口"""
    def __init__(self, **kwargs):
        super(AttentionDecoder, self).__init__(**kwargs)

    @property
    def attention_weights(self):
        raise NotImplementedError

接下来我们在Seq2SeqAttentionDecoder类中实现带有Bahdanau注意力的循环神经网络解码器。首先，初始化解码器的状态，需要下面的输入：

• 编码器在所有时间步的最终层隐状态，作为注意力的键和值；
• 上一时间步的编码器全层隐状态，将作为初始化解码器的隐状态；
• 编码器有效长度（排除在注意力池中填充词元）

在每个解码时间步骤中，解码器上一个时间步的最终层隐状态将用作查询。因此，注意力输出和输入嵌入都连结为循环神经网络解码器的输入。

class Seq2SeqAttentionDecoder(AttentionDecoder):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout=0, **kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        self.attention = AdditiveAttention(num_hiddens, num_hiddens, num_hiddens, dropout)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers, dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)
    
    def init_state(self, enc_outputs, enc_valid_lens, *args):
        # outputs的形状为(batch_size，num_steps，num_hiddens).
        # hidden_state的形状为(num_layers，batch_size，num_hiddens)
        outputs, hidden_state = enc_outputs
        return (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)
    
    def forward(self, X, state):
        # enc_outputs的形状为(batch_size,num_steps,num_hiddens).
        # hidden_state的形状为(num_layers,batch_size,
        # num_hiddens)
        enc_outputs, hidden_state, enc_valid_lens = state
        # 输出x的形状为（num_steps, batch_size, embed_size）
        X = self.embedding(X).permute(1, 0, 2)
        outputs, self._attention_weights = [], []
        for x in X:
            #query的形状为（batch_size, 1, num_hiddens）
            query = torch.unsqueeze(hidden_state[-1], dim=1)
            #context的形状为（batch_size,num_hiddens）
            context = self.attention(query, enc_outputs, enc_outputs, enc_valid_lens)
            #在特征维度上连结
            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
            #将x变形为（1, batch_size,embed_size+num_hiddens）
            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
            outputs.append(out)
            self._attention_weights.append(self.attention.attention_weights)
        #全连接层变换后，outputs的形状为（num_steps, batch_size, vocab_size）
        outputs = self.dense(torch.cat(outputs, dim=0))
        return outputs.permute(1, 0, 2), [enc_outputs, hidden_state, enc_valid_lens]
    
    @property
    def attention_weights(self):
        return self._attention_weights

接下来，使用包含7个时间步的4个序列输入的小批量测试Bahdanau注意力解码器。

class Encoder(nn.Module):
    """编码器-解码器架构的基本编码器接口"""
    def __init__(self, **kwargs):
        super(Encoder, self).__init__(**kwargs)

    def forward(self, X, *args):
        raise NotImplementedError
    
#RNN Encoder
class Seq2SeqEncoder(Encoder):
    """用于序列到序列学习的循环神经网络编码器"""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        #嵌入层
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers, dropout=dropout)

    def forward(self, X, *args):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        X = self.embedding(X)
        # 在循环神经网络模型中，第一个轴对应于时间步
        X = X.permute(1, 0, 2)
        # 如果未提及状态，则默认为0
        output, state = self.rnn(X)
        # output的形状:(num_steps,batch_size,num_hiddens)
        # state的形状:(num_layers,batch_size,num_hiddens)
        return output, state

encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)
encoder.eval()
decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)
decoder.eval()

X = torch.zeros((4, 7), dtype=torch.long)  #(batch_size,num_steps)
state = decoder.init_state(encoder(X), None)
output, state = decoder(X, state)
output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape

输出为:

(torch.Size([4, 7, 10]), 3, torch.Size([4, 7, 16]), 2, torch.Size([4, 16]))

4.3 训练

这里我们指定超参数，实例化一个带有Bahdanau注意力的编码器和解码器，并对这个模型进行机器翻译训练。由于新增的注意力机制，训练要比没有注意力机制的慢得多。

def try_gpu(i=0):
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

from torch.utils import data
import os
import requests
import zipfile
import numpy as np
import matplotlib.pyplot as plt
import time
from IPython import display
import collections

DATA_URL = 'https://d2l-data.s3-accelerate.amazonaws.com/'
DATA_HUB = {'fra-eng': (DATA_URL + 'fra-eng.zip', '94646ad1522d915e7b0f9296181140edcf86a4f5')}

def download_extract(name, folder=None):
    """Download and extract a zip/tar file."""
    url, sha1_hash = DATA_HUB[name]
    fname = os.path.join(folder if folder else '.', url.split('/')[-1])
    
    if os.path.exists(fname):
        print(f"{fname} already exists")
    else:
        print(f"Downloading {fname} from {url}")
        r = requests.get(url, stream=True)
        with open(fname, 'wb') as f:
            f.write(r.content)
        print(f"{fname} downloaded")
    
    base_dir = os.path.splitext(fname)[0]
    if os.path.exists(base_dir):
        print(f"{base_dir} already exists")
    else:
        print(f"Extracting {fname}")
        with zipfile.ZipFile(fname, 'r') as zip_ref:
            zip_ref.extractall(base_dir)
        print(f"Extracted to {base_dir}")
    
    return base_dir

def find_file(root, filename):
    """Recursively find a file in the directory."""
    for dirpath, _, filenames in os.walk(root):
        if filename in filenames:
            return os.path.join(dirpath, filename)
    return None

def read_data_nmt():
    """Load the 'fra-eng' dataset."""
    data_dir = download_extract('fra-eng')
    print(f"Data directory: {data_dir}")
    print(f"Files in data directory: {os.listdir(data_dir)}")
    file_path = find_file(data_dir, 'fra.txt')
    if file_path:
        print(f"Found file at: {file_path}")
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    else:
        raise FileNotFoundError(f"fra.txt not found in directory {data_dir}")
    
def preprocess_nmt(text):
    """预处理"英语－法语"数据集"""
    def no_space(char, prev_char):
        return char in set(',.!?') and prev_char != ' '
    
    # 使用空格替换不间断空格
    # 使用小写字母替换大写字母
    text = text.replace('\u202f', ' ').replace('\xa0', ' ').lower()
    # 在单词和标点符号之间插入空格
    out = [' ' + char if i > 0 and no_space(char, text[i - 1]) else char
           for i, char in enumerate(text)]
    return ''.join(out)

def tokenize_nmt(text, num_examples=None):
    """词元化"英语－法语"数据数据集"""
    source, target = [], []
    for i, line in enumerate(text.split('\n')):
        if num_examples and i > num_examples:
            break
        parts = line.split('\t')
        if len(parts) == 2:
            source.append(parts[0].split(' '))
            target.append(parts[1].split(' '))
    return source, target

raw_text = read_data_nmt()
text = preprocess_nmt(raw_text)

#搭建词表
class Vocab:
    """文本词表"""
    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []
        #按出现频率排序
        counter = count_corpus(tokens)
        self._token_freqs = sorted(counter.items(), key=lambda x: x[1], reverse=True)
        #未知词元的索引维0
        self.idx_to_token = ['<unk>'] + reserved_tokens
        self.token_to_idx = {token: idx
                             for idx, token in enumerate(self.idx_to_token)}
        for token, freq in self._token_freqs:
            if freq < min_freq:
                break
            if token not in self.token_to_idx:
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1

    def __len__(self):
        return len(self.idx_to_token)
    
    def __getitem__(self, tokens):
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, self.unk)
        return [self.__getitem__(token) for token in tokens]

    def to_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]
    
    @property
    def unk(self):#未知词元的索引为0
        return 0
    
    @property
    def token_freqs(self):
        return self._token_freqs
    
source, target = tokenize_nmt(text)

def count_corpus(tokens):
    """统计词元的频率"""
    #这里的tokens是1D列表或2D列表
    if len(tokens) == 0 or isinstance(tokens[0], list):
        #将词元列表展平成一个列表
        tokens = [token for line in tokens for token in line]
    return collections.Counter(tokens)

src_vocab = Vocab(source, min_freq=2,
                      reserved_tokens=['<pad>', '<bos>', '<eos>'])

#截断或填充文本序列
def truncate_pad(line, num_steps, padding_token):
    """截断或填充文本序列"""
    if len(line) > num_steps:
        return line[:num_steps] #截断
    return line + [padding_token] * (num_steps - len(line)) #填充

truncate_pad(src_vocab[source[0]], 10, src_vocab['<pad>'])

#转换文本序列函数
def build_array_nmt(lines, vocab, num_steps):
    """将机器翻译的文本序列转换成小批量"""
    lines = [vocab[l] for l in lines]
    lines = [l + [vocab['<eos>']] for l in lines]
    array = torch.tensor([truncate_pad(
        l, num_steps, vocab['<pad>']) for l in lines])
    valid_len = (array != vocab['<pad>']).type(torch.int32).sum(1)
    return array, valid_len
    
def load_array(data_arrays, batch_size, is_train=True):
    dataset = data.TensorDataset(*data_arrays)
    return data.DataLoader(dataset, batch_size, shuffle=is_train)

def load_data_nmt(batch_size, num_steps, num_examples=600):
    """返回翻译数据集的迭代器和词表"""
    text = preprocess_nmt(read_data_nmt())
    source, target = tokenize_nmt(text, num_examples)
    src_vocab = Vocab(source, min_freq=2, reserved_tokens=['<pad>', '<bos>', '<eos>'])
    tgt_vocab = Vocab(target, min_freq=2, reserved_tokens=['<pad>', '<bos>', '<eos>'])
    src_array, src_valid_len = build_array_nmt(source, src_vocab, num_steps)
    tgt_array, tgt_valid_len = build_array_nmt(target, tgt_vocab, num_steps)
    data_arrays = (src_array, src_valid_len, tgt_array, tgt_valid_len)
    data_iter = load_array(data_arrays, batch_size)
    return data_iter, src_vocab, tgt_vocab

class EncoderDecoder(nn.Module):
    """编码器-解码器架构的基类"""
    def __init__(self, encoder, decoder, **kwargs):
        super(EncoderDecoder, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, enc_X, dec_X, *args):
        enc_outputs = self.encoder(enc_X, *args)
        dec_state = self.decoder.init_state(enc_outputs, *args)
        return self.decoder(dec_X, dec_state)

class Animator:
    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None, ylim=None, xscale='linear', yscale='linear',
                fmts=('-', 'm--', 'g-', 'r:'), nrows=1, ncols=1, figsize=(3.5, 2.5)):
        if legend is None:
            legend = []
            
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes,]

        # self.config_axes = lambda:self.set_axes(
        #     self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)
        self.X, self.Y, self.fmts = None, None, fmts
        self.set_axes(xlabel, ylabel, xlim, ylim, xscale, yscale, legend)

    def set_axes(self, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):
        """设置轴"""
        for ax in self.axes:
            ax.set_xlabel(xlabel)
            ax.set_ylabel(ylabel)
            ax.set_xscale(xscale)
            ax.set_yscale(yscale)
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)
            if legend:
                ax.legend(legend)
            ax.grid()
            
    def add(self, x, y):
        if not hasattr(y, '__len__'):
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)

        for ax in self.axes:
            ax.cla()
            for x, y, fmt in zip(self.X, self.Y, self.fmts):
                ax.plot(x, y, fmt)
        display.display(self.fig)
        display.clear_output(wait=True)
        
#定义计时器
class Timer:
    def __init__(self):
        self.times = []
        self.start()

    def start(self):
        self.tik = time.time()

    def stop(self):
        self.times.append(time.time() - self.tik)
        return self.times[-1]
    
    def avg(self):
        return sum(self.times) / len(self.times)
    
    def sum(self):
        return sum(self.times)
    
    def cumsum(self):
        return np.array(self.times).cumsum().tolist()
    
#定义程序Accumulator
class Accumulator:  #@save
    """在n个变量上累加"""
    def __init__(self, n):
        self.data = [0.0] * n

    def add(self, *args):
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def reset(self):
        self.data = [0.0] * len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]
    
#裁剪梯度
def grad_clipping(net, theta):
    """裁剪梯度"""
    if isinstance(net, nn.Module):
        params = [p for p in net.parameters() if p.requires_grad]
    else:
        params = net.params
    norm = torch.sqrt(sum(torch.sum((p.grad ** 2)) for p in params))
    if norm > theta:
        for param in params:
            param.grad[:] *= theta / norm


class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""
    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction='none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss

def train_seq2seq(net, data_iter, lr, num_epochs, tgt_vocab, device):
    """训练序列到序列模型"""
    def xavier_init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        if type(m) == nn.GRU:
            for param in m._flat_weights_names:
                if 'weight' in param:
                    nn.init.xavier_uniform_(m._parameters[param])

    net.apply(xavier_init_weights)
    net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    loss = MaskedSoftmaxCELoss()
    net.train()
    animator = Animator(xlabel='epoch', ylabel='loss', xlim=[10, num_epochs])
    for epoch in range(num_epochs):
        timer = Timer()
        metric = Accumulator(2)
        for batch in data_iter:
            optimizer.zero_grad()
            X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]
            bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0], device=device).reshape(-1,1)
            dec_input = torch.cat([bos, Y[:, :-1]], 1) #强制教学
            Y_hat, _ = net(X, dec_input, X_valid_len)
            l = loss(Y_hat, Y, Y_valid_len)
            l.sum().backward()
            grad_clipping(net, 1)
            num_tokens = Y_valid_len.sum()
            optimizer.step()
            with torch.no_grad():
                metric.add(l.sum(), num_tokens)
        if (epoch + 1) % 10 == 0:
            animator.add(epoch + 1, (metric[0] / metric[1], ))
    print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} 'f'tokens/sec on {str(device)}')

embed_size, num_hiddens, num_layers, dropout =32, 32, 2, 0.1
batch_size, num_steps = 64, 10
lr, num_epochs, device = 0.005, 250, try_gpu()

train_iter, src_vocab, tgt_vocab = load_data_nmt(batch_size, num_steps)
encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hiddens, num_layers, dropout)
decoder = Seq2SeqAttentionDecoder(len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
net = EncoderDecoder(encoder, decoder)
train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

可以得到结果为：

loss 0.020, 6838.0 tokens/sec on cuda:0

模型训练后，我们用它将几个英语句子翻译成法语并计算它们的BLEU分数。

#预测
def predict_seq2seq(net, src_sentence, src_vocab, tgt_vocab, num_steps,
                    device, save_attention_weights=False):
    """序列到序列模型的预测"""
    # 在预测时将net设置为评估模式
    net.eval()
    src_tokens = src_vocab[src_sentence.lower().split(' ')] + [
        src_vocab['<eos>']]
    enc_valid_len = torch.tensor([len(src_tokens)], device=device)
    src_tokens = truncate_pad(src_tokens, num_steps, src_vocab['<pad>'])
    # 添加批量轴
    enc_X = torch.unsqueeze(
        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)
    enc_outputs = net.encoder(enc_X, enc_valid_len)
    dec_state = net.decoder.init_state(enc_outputs, enc_valid_len)
    # 添加批量轴
    dec_X = torch.unsqueeze(torch.tensor(
        [tgt_vocab['<bos>']], dtype=torch.long, device=device), dim=0)
    output_seq, attention_weight_seq = [], []
    for _ in range(num_steps):
        Y, dec_state = net.decoder(dec_X, dec_state)
        # 我们使用具有预测最高可能性的词元，作为解码器在下一时间步的输入
        dec_X = Y.argmax(dim=2)
        pred = dec_X.squeeze(dim=0).type(torch.int32).item()
        # 保存注意力权重（稍后讨论）
        if save_attention_weights:
            attention_weight_seq.append(net.decoder.attention_weights)
        # 一旦序列结束词元被预测，输出序列的生成就完成了
        if pred == tgt_vocab['<eos>']:
            break
        output_seq.append(pred)
    return ' '.join(tgt_vocab.to_tokens(output_seq)), attention_weight_seq

#BLEU实现
def bleu(pred_seq, label_seq, k): 
    """计算BLEU"""
    pred_tokens, label_tokens = pred_seq.split(' '), label_seq.split(' ')
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    score = math.exp(min(0, 1 - len_label / len_pred))
    for n in range(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        for i in range(len_label - n + 1):
            label_subs[' '.join(label_tokens[i: i + n])] += 1
        for i in range(len_pred - n + 1):
            if label_subs[' '.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[' '.join(pred_tokens[i: i + n])] -= 1
        score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    return score

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 = predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {bleu(translation, fra, k=2):.3f}')

输出有：

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

训练结束后，下面通过可视化注意力权重 会发现，每个查询都会在键值对上分配不同的权重，这说明 在每个解码步中，输入序列的不同部分被选择性地聚集在注意力池中。

attention_weights = torch.cat([step[0][0][0] for step in dec_attention_weight_seq], 0).reshape((
    1, 1, -1, num_steps))

# 加上一个包含序列结束词元
show_heatmaps(
    attention_weights[:, :, :, :len(engs[-1].split()) + 1].cpu(),
    xlabel='Key positions', ylabel='Query positions')

5.多头注意力

在实践中，当给定相同的查询、键和值的集合时，我们希望模型可以基于相同的注意力机制学习到不同的行为，然后将不同的行为作为知识组合起来，捕获序列内各种范围的依赖关系（例如，短距离依赖和长距离依赖关系）。因此，允许注意力机制组合使用查询、键和值的不同子空间表示（representation subspaces）可能是有益的。

为此，与其只使用单独一个注意力汇聚，可以使用独立学习得到的h组不同的线性投影（linear projections）来变换查询、键和值。然后这h组变换后的查询、键和值将并行地送到注意力汇聚中。最后，将这h个注意力汇聚的输出拼接在一起，并且通过另一个可以学习的线性投影进行变换，以产出最终输出。这种设计被称为多头注意力（multihead attention）。对于h个注意力汇聚输出，每一个注意力汇聚都被称作一个头（head）。下图展示了使用全连接层来实现可学习的线性变换的多头注意力。

5.1 模型

对于多头注意力，在数学上可以有以下表达。给定查询 q ∈ R d q \mathbf{q} \in \mathbb{R}^{d_q} q∈Rdq、键 k ∈ R d k \mathbf{k} \in \mathbb{R}^{d_k} k∈Rdk和值 v ∈ R d v \mathbf{v} \in \mathbb{R}^{d_v} v∈Rdv，每个注意头 h i ( i = 1 , ... , h ) \mathbf{h}i(i=1, \dots, h) hi(i=1,...,h)的计算方法为：
h i = f ( W i ( q ) q , W i ( k ) k , W i ( v ) v ) ∈ R p v \mathbf{h}{i}=f\left(\mathbf{W}{i}^{(q)} \mathbf{q}, \mathbf{W}{i}^{(k)} \mathbf{k}, \mathbf{W}{i}^{(v)} \mathbf{v}\right) \in \mathbb{R}^{p{v}} hi=f(Wi(q)q,Wi(k)k,Wi(v)v)∈Rpv

其中，可学习的参数包括 W i ( q ) ∈ R p q × d q \mathbf{W}{i}^{(q)} \in \mathbb{R}^{p_q \times d_q} Wi(q)∈Rpq×dq、 W i ( k ) ∈ R p k × d k \mathbf{W}{i}^{(k)} \in \mathbb{R}^{p_k \times d_k} Wi(k)∈Rpk×dk和 W i ( v ) ∈ R p v × d v \mathbf{W}{i}^{(v)} \in \mathbb{R}^{p_v \times d_v} Wi(v)∈Rpv×dv，以及代表注意力汇聚的函数f。f是可以是加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换，他对应着h个头连结后的结果，因此其可学习参数是 W Q ∈ R p o × h p \mathbf{W}{Q} \in \mathbb{R}^{p_o \times h_p} WQ∈Rpo×hp:
W o ∣   h 1 ⋮ h h ] ] ∈ R p o \left.\begin{array}{c} \mathbf{W}{o} \left.\left\lvert\, \begin{array}{c} \mathbf{h}{1} \\ \vdots \\ \mathbf{h}{h} \end{array}\right.\right] \end{array}\right] \in \mathbb{R}^{p{o}} Wo h1⋮hh ∈Rpo

基于这种设计，每个头都可能会关注输入的不同部分，可以表示比简单加权平均值更复杂的函数。

5.2 实现

在实现过程中通常选择缩放点积注意力作为每个注意力头。为了避免计算代价和参数代价的大幅增长，我们设定 p q = p k = p v = p o / h o p_q=p_k=p_v=p_o/h_o pq=pk=pv=po/ho。值得注意的是，如果将查询、键和值的线性变换的输出数量设置为 p q h = p k h = p v h = p o p_qh=p_kh=p_vh=p_o pqh=pkh=pvh=po，则可以并行计算h个头。接下来对此进行实现， p o p_o po通过参数num_hiddens指定。

import math
import torch 
from torch import nn

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)

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

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

输出有：

MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)

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

输出为：

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

6.自注意力和位置编码

在深度学习中，经常使用CNN或RNN对序列进行编码。有了注意力机制之后，我们将词元序列输入注意力池化中，以便同一组词元同时充当查询、键和值。具体来说，每个查询都会关注多有的键值对并生成一个注意力输出。由于查询、键和值来自同一组输入，因此被称为自注意力（self-attention），也被称为内部注意力（intra-attention）。

6.1 自注意力

给定一个由词元组成的输入序列 x 1 , ... , x n \mathbf{x}_1, \dots, \mathbf{x}_n x1,...,xn，其中任意 x i ∈ R d ( 1 ≤ i ≤ n ) \mathbf{x}_i \in \mathbb{R}^d(1 \le i \le n) xi∈Rd(1≤i≤n)。该序列的自注意力输出为一个长度相同的序列 y 1 , ... , y n \mathbf{y}_1, \dots, \mathbf{y}_n y1,...,yn，其中：
y i = f ( x i , ( x 1 , x 1 ) , ... , ( x n , x n ) ) ∈ R d \mathbf{y}_i = f(\mathbf{x}_i,(\mathbf{x}_1, \mathbf{x}_1), \dots, (\mathbf{x}_n, \mathbf{x}_n)) \in \mathbb{R}^d yi=f(xi,(x1,x1),...,(xn,xn))∈Rd

根据之前定义的注意力汇聚函数f，下面的代码基于多头注意力对一个张量完成自注意力的计算，张量的形状为（批量大小，时间步的数目或词元序列的长度，d）。输出与输入的张量形状相同。

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

输出有：

MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)

同样查看形状：

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

有输出：

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

6.2 卷积神经网络、循环神经网络和自注意力的比较

对于这三个架构，目标都是将n个词元组成的序列映射到另一个长度相等的序列，其中的每个输入词元或输出词元都由d维向量表示。具体来说，将比较的是卷积神经网络、循环神经网络和自注意力这几个架构的计算复杂性、顺序操作和最大路径长度。需要注意的是，顺序操作会妨碍并行计算，而任意的序列位置组合之间的路径越短，则能更轻松地学习序列中的远距离依赖关系。

考虑一个卷积核大小为k的卷积层，由于序列长度为n，输入和输出的通道数量都是d，所以卷积层的计算复杂度 O ( k n d 2 ) \mathcal{O} (knd^2) O(knd2)，如上图所示，卷积神经网络是分层的，因此为有 O ( 1 ) \mathcal{O} (1) O(1)个顺序操作，最大路径长度为 O ( n / k ) \mathcal{O} (n/k) O(n/k)。例如， x 1 \mathbf{x}_1 x1和 x 5 \mathbf{x}_5 x5处于图中卷积核大小为3的双层卷积神经网络的感受野内。

当更新循环神经网络的隐状态时， d × d d \times d d×d权重矩阵和d维隐状态的乘法计算复杂度为 O ( d 2 ) \mathcal{O} (d^2) O(d2)。由于序列长度为n，因此循环神经网络层的计算复杂度为 O ( n d 2 ) \mathcal{O} (nd^2) O(nd2)。由图可知，有 O ( n ) \mathcal{O} (n) O(n)个顺序操作无法并行化，最大路径长度也是 O ( n ) \mathcal{O} (n) O(n).

在子注意力中，查询、键和值都是 n × d n \times d n×d矩阵。考虑缩放点积注意力，其中 n × d n \times d n×d矩阵乘以 d × n d \times n d×n矩阵，之后输出的 n × n n \times n n×n矩阵乘以 n × d n \times d n×d矩阵。因此，自注意力具有 O ( n 2 d ) \mathcal{O} (n^2d) O(n2d)计算复杂度。如图，每个词元都通过自注意力直接连接到任何其他词元，因此有 O ( 1 ) \mathcal{O} (1) O(1)个顺序操作可以并行计算，最大路径长度也是 O ( 1 ) \mathcal{O} (1) O(1)。

总而言之，CNN和子注意力都拥有并行计算的优势，而且自注意力的最大路径长度最短。但是他在很长的序列中计算会很慢。

6.3 位置编码

在处理词元序列时，循环神经网络是逐个重复地处理词元地，而自注意力则因为并行计算而放弃了顺序操作。为了使用序列的顺序信息，通过在输入表示中添加位置编码（positional encoding）来注入绝对的或相对的位置信息。位置编码可以通过学习得到也可以直接固定得到。接下来描述的是基于正弦函数和余弦函数的固定位置编码。

假设输入表示 X ∈ R n × d \mathbf{X} \in \mathbb{R}^{n \times d} X∈Rn×d包含一个序列中n个词元的d维嵌入表示。位置编码使用相同形状的位置嵌入矩阵 P ∈ R n × d \mathbf{P} \in \mathbb{R}^{n \times d} P∈Rn×d输出 X + P \mathbf{X + P} X+P，矩阵第i行、第2j列和第2j+1列的元素为：
p i , 2 j = sin ⁡ ( i 1000 0 2 j / d ) , p i , 2 j + 1 = cos ⁡ ( i 1000 0 2 j / d ) . \begin{aligned} p_{i, 2 j} & =\sin \left(\frac{i}{10000^{2 j / d}}\right), \\ p_{i, 2 j+1} & =\cos \left(\frac{i}{10000^{2 j / d}}\right) . \end{aligned} pi,2jpi,2j+1=sin(100002j/di),=cos(100002j/di).

我们定义PositionalEncoding类来实现它。

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)

在位置嵌入矩阵 P \mathbf{P} P中，行代表词元在序列中的位置，列代表位置编码的不同维度。从下面的例子可以看出位置嵌入矩阵的第6列和第7列的频率高于第8列和第9列。第6列和第7列之间的偏移量（第8列和第9列相同）是由于正弦函数和余弦函数的交替。

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]

plt.figure(figsize=(6, 2.5))
for d in range(6, 10):
    plt.plot(torch.arange(num_steps), P[0, :, d].numpy(), label=f"Col {d}")

plt.xlabel('Row (position)')
plt.legend()
plt.show()

6.3.1 绝对位置信息

为了明白沿着编码维度单调降低的频率与绝对位置信息的关系，我们打印出0-7的二进制表现形式。正如所看到的，每个数字、没两个数字和每四个数字的比特值在第一个最低位、第二个最低位和第三个最低位上分别交替。

for i in range(8):
    print(f'{i}的二进制是：{i:>03b}')

在二进制表现中，较高比特位的交替频率低于较低比特位，与下面的热图所示相似，只是位置编码通过使用三角函数在编码维度上降低频率。由于输出是浮点数，因此此类连续表示比二进制表示法更节省空间。

P = P[0, :, :].unsqueeze(0).unsqueeze(0)
show_heatmaps(P, xlabel='Column (encoding dimension)',
                  ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

6.3.2 相对位置信息

除了捕获绝对信息之外，上述的位置编码还允许模型学习得到输入序列中相对位置信息。这是因为对于任何确定的位置偏移 δ \delta δ，位置 i + δ i+\delta i+δ处的位置编码可以线性投影位置i处的位置编码来表示。

这种投影的数学解释是，令 ω j = 1 / 1000 0 2 j / d \omega_j = 1/10000^{2j/d} ωj=1/100002j/d，对于任何确定的位置偏移 δ \delta δ ,位置编码中的任何一对 ( p i , 2 j , p i , 2 j + 1 ) (p_{i,2j}, p_{i,2j+1}) (pi,2j,pi,2j+1)都可以线性投影到 ( p i + δ , 2 j , p i + δ , 2 j + 1 ) (p_{i+\delta ,2j},p_{i+ \delta ,2j+1}) (pi+δ,2j,pi+δ,2j+1)：
[ cos ⁡ ( δ ω j ) sin ⁡ ( δ ω j ) − sin ⁡ ( δ ω j ) cos ⁡ ( δ ω j ) ] [ p i , 2 j p i , 2 j + 1 ] = [ cos ⁡ ( δ ω j ) sin ⁡ ( i ω j ) + sin ⁡ ( δ ω j ) cos ⁡ ( i ω j ) − sin ⁡ ( δ ω j ) sin ⁡ ( i ω j ) + cos ⁡ ( δ ω j ) cos ⁡ ( i ω j ) ] = [ sin ⁡ ( ( i + δ ) ω j ) cos ⁡ ( ( i + δ ) ω j ) ] = [ p i + δ , 2 j p i + δ , 2 j + 1 ] \begin{aligned} & {\left[\begin{array}{cc} \cos \left(\delta \omega_{j}\right) & \sin \left(\delta \omega_{j}\right) \\ -\sin \left(\delta \omega_{j}\right) & \cos \left(\delta \omega_{j}\right) \end{array}\right]\left[\begin{array}{c} p_{i, 2 j} \\ p_{i, 2 j+1} \end{array}\right] } \\ = & {\left[\begin{array}{c} \cos \left(\delta \omega_{j}\right) \sin \left(i \omega_{j}\right)+\sin \left(\delta \omega_{j}\right) \cos \left(i \omega_{j}\right) \\ -\sin \left(\delta \omega_{j}\right) \sin \left(i \omega_{j}\right)+\cos \left(\delta \omega_{j}\right) \cos \left(i \omega_{j}\right) \end{array}\right] } \\ = & {\left[\begin{array}{c} \sin \left((i+\delta) \omega_{j}\right) \\ \cos \left((i+\delta) \omega_{j}\right) \end{array}\right] } \\ = & {\left[\begin{array}{c} p_{i+\delta, 2 j} \\ p_{i+\delta, 2 j+1} \end{array}\right] } \end{aligned} ===[cos(δωj)−sin(δωj)sin(δωj)cos(δωj)][pi,2jpi,2j+1][cos(δωj)sin(iωj)+sin(δωj)cos(iωj)−sin(δωj)sin(iωj)+cos(δωj)cos(iωj)][sin((i+δ)ωj)cos((i+δ)ωj)][pi+δ,2jpi+δ,2j+1]
2 × 2 2 \times 2 2×2的投影矩阵不依赖于任何位置的索引i。

7.Transformer

7.1 模型

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

从宏观角度来看，Transformer的编码器是由多个相同的层叠加而成的，每个层都有两个子层（sublayer）。第一个子层是多头自注意力（multi-head self-attention）汇聚；第二个子层是基于位置的前馈网络（positionwise feed-forward network）。具体来说，在计算编码器的自注意力时，查询、键和值都来自前一个编码器层的输出。受残差网络的启发，每个子层都采用了残差连接（residual connection）。在Transformer中，对于序列中任何位置输入 x ∈ R d \mathbf{x} \in \mathbb{R}^d x∈Rd都要求满足 s u b l a y e r ( x ) ∈ R d sublayer(\mathbf{x}) \in \mathbb{R}^d sublayer(x)∈Rd，以便残差连接满足 x + s u b l a y e r ( x ) ∈ R d \mathbf{x} + sublayer(\mathbf{x}) \in \mathbb{R}^d x+sublayer(x)∈Rd。在残差连接的加法计算之后，紧接着应用层规范化（layer normalization）。因此，输入序列对应的每个位置，Transformer编码器都将输出一个d维表示向量。

Transformer解码器也是由多个相同的层叠加而成的，并且层中使用了残差连接和层规范化。除了编码器中描述的两个子层之外，解码器还在这两个子层之间插入了第三个子层，称为编码器-解码器注意力层（encoder-decoder attention layer）。在编码器-解码器注意力中，查询、键和值都来自上一个解码器层的输出，而键和值来自整个编码器的输出。但是解码器中的每个位置只能考虑该位置之前的所有位置。这种掩蔽（masked）注意力保留了自回归（auto-regressive）属性，确保预测仅依赖于已生成的输出词元。

在此之前已经实现了基于缩放点积多头注意力和位置编码，接下来实现Transformer模型的剩余部分。

7.2 基于位置的前馈网络

基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的原因。接下来我们实现输入X的形状（批量大小，时间步数或序列长度，隐单元数或特征维度）将被一个两层的感知机转换成形状为（批量大小，时间步数，ffn_num_outputs）的输出张量。

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

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.1504,  0.2580, -0.4174, -0.0032,  0.5616, -0.0662,  0.3975,  0.1379],
        [-0.1504,  0.2580, -0.4174, -0.0032,  0.5616, -0.0662,  0.3975,  0.1379],
        [-0.1504,  0.2580, -0.4174, -0.0032,  0.5616, -0.0662,  0.3975,  0.1379]],
       grad_fn=<SelectBackward0>)

7.3 残差连接和层规范化

层规范化和批量规范化的目标相同，但是层规范化时基于特征维度进行规范化，尽管批量规范化在计算机视觉中被广泛应用，但在自然语言处理任务中（输入通常是变长序列）批量规范化通常不如层规范化的效果好。

接下来使用残差连接和层规范化来实现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)

需要注意的是，残差连接要求两个输入的形状相同，以便加法操作后输出张量的形状相同。

7.4 编码器

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

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 = 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编码器中的任何层都不会改变其输入的形状。

接下来实现堆叠了num_layers个EncoderBlock的Transformer编码器。由于这里使用的是值范围在-1到1之间的固定位置编码，因此通过学习得到的输入的嵌入表示的值需要先乘以嵌入维度的平方根进行重新缩放，然后再与位置编码相加。

class TransformerEncoder(nn.Module):
class TransformerEncoder(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 = 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编码器。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

输出为：

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

7.5 解码器

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

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

class DecoderBlock(nn.Module):
    """解码器块"""
    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

        #自注意力
        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)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

为了便于在编码器-解码器注意力中进行缩放点积计算和残差连接中进行加法计算，编码器和解码器的特征维度都是num_hiddens。

现在我们构建由num_layers个DecoderBlock实例组成的完整的Transformer解码器。最后，通过一个全连接层计算所有vocab_size个可能的输出词元的预测值。解码器的自注意力权重和编码器解码器注意力权重都被存储下来，方便日后可视化的需要。

class TransformerDecoder(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):
        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 = 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

7.6 训练

在这里，指定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, 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 = 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 = EncoderDecoder(encoder, decoder)
train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

输出有：

loss 0.033, 5306.8 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 = predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {bleu(translation, fra, k=2):.3f}')

输出为：

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

当进行最后一个英语到法语的句子翻译工作时，让我们可视化Transformer的注意力权重。编码器自注意力权重的形状为（编码器层数，注意力头数，num_steps或查询的数目，num_steps或"键－值"对的数目）。

enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
    -1, num_steps))
enc_attention_weights.shape

输出为：

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

在编码器的自注意力中，查询和键都来自相同的输入序列。因为填充词元是不携带信息的，因此通过指定输入序列的有效长度可以避免查询与使用填充词元的位置计算注意力。接下来，将逐行呈现两层多头注意力的权重。每个注意力头都根据查询、键和值的不同的表示子空间来表示不同的注意力。

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

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

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=(9, 3.5))

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

show_heatmaps(
    dec_inter_attention_weights, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(9, 3.5))

尽管Transformer架构是为了序列到序列的学习而提出的，但Transformer编码器或Transformer解码器通常被单独用于不同的深度学习任务中。