当人类看到一张图片或者听到一段文字时，人类不会将其中每个像素或者每个字都细致观察，而是会选择其中较为重要的部分进行处理。例如对于以下图片，我们第一眼会注意到图片中的长颈鹿，因为它们处于图片中央，且与背景相差较大。也有可能你第一眼注意到了后面的树或者地面，不过不论如何，都有个首先关注的重点。接下来，假设我们需要猜出这张图片拍摄的场所，那么我们会关注到背景中的栅栏，从而推测出拍摄于动物园。在同一时间，人眼不会将注意力放到所有的物体上，而是会有最关注的物体，例如你读这句话的时候焦点就在这里，而不是其他地方，这就是Attention。

图长颈鹿

这种机制可以帮助我们更快、更准确地获取和处理信息。它的核心思想是基于原有的数据找到关联性，然后突出其某些重要的特征。

Seq2Seq模型

在介绍怎么计算Attention之前，先简要介绍一下Seq2Seq模型，全称Sequence to Sequence，是以前机器翻译等自然语言处理任务常用的模型。

图 Seq2Seq模型

首先，当我们处理一段文字时，我们通常会将它划分成一个个文本数据的最小单元或基本单元，称之为token。划分token的方式很多，包括字符划分、单词划分等。用以下一句话举例：

我喜欢小狗。 | I like dogs.

在这个例子中，我们可以将其划分为六个中文token："我"、"喜"、"欢"、"小"、"狗"、"。"和四个英文token："I"、"like"、"dogs"、"."。在Seq2Seq模型中，为了增加上下文信息以更好地理解输入和输出序列的结构，在每个句子的开头和结尾添加了特殊的<SOS>（Start of Sentence）和<EOS>（End of Sentence）两个token。

Seq2Seq模型采用了Encoder-Decoder（编码器-解码器）架构。该模型的基本操作包括将一段token按顺序输入编码器，获得有关内容的编码；然后，将该编码与<SOS>一同输入解码器，生成第一个预测token，随后将该token作为下一个输入，循环此过程直到输出的token为<EOS>为止。

Attention的计算

考虑上述例子，我们很容易想到，在翻译过程中，"dogs"与输入的中文token "狗"存在强相关性，与"小"也具有相关性，而与其他token相关性较弱。如果转换为注意力机制，这意味着对每个token的注意力分数不同。要实现这个功能，则需要给解码器配备某种搜索功能，允许它在需要生成输出词时查看整个源句子，即注意力机制。

理论上，我们已经了解了注意力机制的运作方式，现在是时候从技术层面深入理解了。我们依旧以刚才的“我喜欢小狗。”中输出“dogs”这一步为例。计算步骤如下：

计算每个编码器状态的得分：假设在Seq2Seq的编码器中，输入产生的隐变量分别为h1、h2、h3、h4、h5、h6，训练一个神经网络计算注意力分，其对应的注意力分为s1、s2、s3、s4、s5、s6。其中s4和s5相对较高，其他相对较低。
计算注意力权重：生成分数后，对这些分数使用softmax 以获得权重 w1、w2、w3、w4、w5、w6。
计算上下文向量：上下文向量c= w1*s1+w2*s2+w3*s3+w4*s4+w5*s5+w6*s6 。
将上下文向量与解码器之前的输出内容（即“like”）结合共同输入模型。

以下两张图可以很明显得对比加入Attention前后的Seq2Seq模型变化：

图加入Attention前后的Seq2Seq模型

Transformer与自注意力机制

Transformer是一种用于处理序列数据的深度学习模型，广泛用于自然语言处理任务。它于2017年由Google Brain的研究人员Vaswani等人在论文《Attention is All You Need》中首次提出。

在该模型中，Transformer提出了自注意力机制（Self-Attention），通过这种机制，模型能够同时关注序列中的所有位置，而不是依赖于传统的递归或卷积结构。这种全局性的关注机制使得Transformer能够在处理长距离依赖性的任务上取得显著的性能提升。

相比于无论在训练还是测试过程中都强依赖上一时刻输出的结果的循环神经网络及其变体，Transformer的设计强调了并行计算的能力，使其在训练中能够并行计算出所需的序列从而避免循环输入。所有token是同时训练的，这样大大增加了计算效率。。

由于其出色的性能和可扩展性，Transformer极大地推动了自然语言处理领域的发展，成为许多先进模型（如BERT、GPT等）乃至现在大语言模型的基础，后续至今的许多AI领域都引入了Transformer模型。 Transformer的出现标志着深度学习在处理序列数据方面取得的一项重大突破。

Transformer的结构图如下，接下来我会结合代码对每个部分进行详细分析，并解释它这么设计的原因。

图 Transformer架构

总体架构

Transformer 模型也使用Encoder-Decoder架构，整体结构分为两个主要组件：Encoder（编码器）和Decoder（解码器），这两部分共同协作完成对输入序列到输出序列的映射。这一创新性的结构不同于传统的递归或卷积结构，而是引入了自注意力机制，极大地提高了模型在处理序列数据时的效率和性能。

首先，让我们聚焦于Encoder。Encoder负责将输入序列（例如语言文本）映射为隐藏层表示。在具体实现中，Encoder由多个相同的Encoder Layer（编码器层）组成，每个Encoder Layer内部又包含两个主要子层：Multi-Head Self-Attention（多头自注意力机制）和 Feedforward Neural Network（前馈神经网络）。这一层层的结构有助于模型更好地捕捉输入序列的语义和特征。

接下来是Decoder。Decoder同样由多个Decoder Layer组成，每个Decoder Layer内部包含三个子层：Masked Multi-Head Self-Attention（掩码多头自注意力机制）、Multi-Head Self-Attention和Feedforward Neural Network。Decoder的目标是将隐藏层映射为目标语言序列。在此过程中，模型通过掩码机制确保在生成每个标记时只能依赖于先前生成的标记，以避免信息泄漏。

过去没有解决不同token间如何有效相互影响的问题，而Transformer关键的创新点在于自注意力机制，它允许模型在处理输入序列时动态地关注不同位置的信息，而不受限于固定的感受野。这使得Transformer能够同时处理长距离的依赖性，有力地解决了传统模型在此方面的局限性。

位置编码(Positional Encoding)

介绍

后面分析自注意力机制时可以看到，Transformer中的每一个部分计算都是与位置无关的，而且由于是并行计算，因此输入并没有时域区分。基于这个原因，需要一种机制来让模型知道输入序列中各个标记的相对位置。

位置编码的设计是通过在输入嵌入中添加一些额外的信息来实现的。位置嵌入的维度为 [max_sequence_length, embedding_dimension], 位置嵌入的维度与词向量的维度是相同的，都是 embedding_dimension。max_sequence_length 属于超参数，指的是限定每个句子最长由多少个词构成。位置编码会在Word Embedding后直接与其在数值上相加。

在论文《Attention is All You Need》中，使用的是正余弦函数位置编码，公式如下：

其中，pos 指的是一句话中某个字的位置，取值范围是[0, max_sequence_length)，i指的是embedding的第i个维度，d 指的是embedding_dimension。

这一组公式对应着每个嵌入维度的一组奇数和偶数的序号，例如，0和 1 一组，2 和 3 一组，以此类推。这些公式利用上述的正弦和余弦函数处理，每个位置在 embedding_dimension 维度上都会得到不同周期的正弦和余弦函数的取值组合。这样，就能够在嵌入中产生独特的纹理，捕捉到位置信息的变化。随着维度序号的增大，正弦和余弦函数的周期变化越来越慢，最终形成一种包含位置信息的纹理，使得模型能够学到位置之间的依赖关系和自然语言的时序特性。

以上一段话肯定会让人看着云里雾里，那么我们尝试从设计角度理解。理想情况下，位置编码的设计应该满足以下条件：

它应该为每个位置输出唯一的编码。
不同长度的句子之间，任何两个字之间的差值应该保持一致。
它的值应该是有界的。

位置编码不是一个数字，而是一个包含句子中特定位置信息的embedding_dimension维向量。这个向量是用来给句子中的每个字提供位置信息的，换句话说，我们通过注入每个字位置信息的方式，增强了模型的输入。正如上面所说，是将位置嵌入和字嵌入相加，然后作为输入。

通过绘制的位置嵌入图，纵向观察，可见随着 embedding_dimension序号增大，位置嵌入函数的周期变化越来越平缓。其中每一行都是一个位置的位置嵌入，不同位置之间的位置嵌入不同，可以想象成是二进制编码的形式。

图位置编码可视化

除了这种提前计算的位置编码，如果输入长度是固定的，也可以考虑使用可学习的位置编码，效果也还可以。

代码

可学习位置编码：

class LearnablePositionalEncoding(nn.Module):

# Learnable positional encoding

def __init__(self, emb_dim, len):

super(LearnablePositionalEncoding, self).__init__()

assert emb_dim > 0 and len > 0, 'emb_dim and len must be positive'

self.emb_dim = emb_dim

self.len = len

self.pe = nn.Parameter(torch.zeros(len, emb_dim))

def forward(self, x):

return x + self.pe[:x.size(-2), :]

正余弦位置编码：

class PositionalEncoding(nn.Module):

# Sine-cosine positional coding

def __init__(self, emb_dim, max_len, freq=10000.0):

super(PositionalEncoding, self).__init__()

assert emb_dim > 0 and max_len > 0, 'emb_dim and max_len must be positive'

self.emb_dim = emb_dim

self.max_len = max_len

self.pe = torch.zeros(max_len, emb_dim)

pos = torch.arange(0, max_len).unsqueeze(1)

# pos: [max_len, 1]

div = torch.pow(freq, torch.arange(0, emb_dim, 2) / emb_dim)

# div: [ceil(emb_dim / 2)]

self.pe[:, 0::2] = torch.sin(pos / div)

# torch.sin(pos / div): [max_len, ceil(emb_dim / 2)]

self.pe[:, 1::2] = torch.cos(pos / (div if emb_dim % 2 == 0 else div[:-1]))

# torch.cos(pos / div): [max_len, floor(emb_dim / 2)]

多头自注意力(Multi-Head Attention)

自注意力(Self Attention)

我们输入的句子维度是[max_sequence_length]，经过词嵌入和位置编码之后，变成了[max_sequence_length, embedding_dimension]，其中embedding_dimension这个维度是对这个词的一个表达。

我们先直观得思考一下，如果我们需要知道某个词xi 和其他所有词之间会怎么相互影响，并且相互影响之后怎么得到结果。那么直观得想，我们使用一个qi 来指代xi ，让这个qi 作为一个询问，称之为查询(queue)。同时，让所有词都给出一个自己待查询的指代，写作k ，称为键(key)。那么我们将qi 与每个词的k 进行匹配，得出一个匹配系数，或者可以理解成相关系数。之后我们就知道xi 与每个词（包括它自己）之间的相关程度。由于匹配系数未经过约束，可能极大或者极小，那么我们通过一个Softmax函数对其进行约束，使其归一化。最后，我们再让每个词给出自己的内容的一个指代v ，称之为值(value)。让每个词与这个词xi 的相关系数乘以这个词的值，再全部加起来，就得到了更新后的这个词xi。

详细的计算步骤如下：

输入表示：假设我们有一个输入序列X= x1, x2… xn ，其中n 表示序列长度。对于每个元素xi ，我们将其映射成三个向量，即查询向量Qi 、键向量Ki 和数值向量Vi ，这三个向量是通过对xi 应用线性映射得到的。

其中，WQ ，WK ，WV 分别是查询、键和数值的线性映射矩阵。

图计算Q，K，V

注意力分数：通过计算查询向量和键向量的点积，然后除以缩放系数（其中dk 是查询向量或键向量的维度），得到注意力分数（Attention Scores）。

图计算注意力分数

注意力权重：将注意力分数通过Softmax函数进行归一化，得到注意力权重。这些权重表示了每个位置对于当前位置的相对重要性。

图注意力分数Softmax

加权和：使用注意力权重对数值向量进行加权和，得到自注意力的输出。

图数值加权

图得到输出结果

矩阵计算

上面介绍的图中方法需要一个循环遍历所有的token，我们可以把上面的向量计算变成矩阵的形式，从而一次计算出所有token的输出，第一步就不是计算某个token的Qi 了，而是一次计算所有token的Q ，K 和V 。计算过程如下图所示，这里的输入是一个矩阵X ，矩阵第i行就是xi 。

图矩阵计算Q，K，V

接下来将Q 和K转置相乘，然后除以（为了让梯度正常），经过 softmax 以后再乘以V 得到输出。

图矩阵得到输出结果

多头(Multi-head)

多头的理解相对简单。我们前面定义的一组Q ，K ，V 可以让一个词关乎到相关的内容，我们可以定义多组Q ，K ，V ，让它们分别关注不同的上下文。计算的过程还是一样，只不过线性变换的矩阵从一组线性映射矩阵变成了多组线性映射矩阵，如下图所示：

图多头Q，K，V

对于输入矩阵X ，每组Q ，K ，V 都能得到一个输出矩阵，最后我们将所有输出矩阵合起来拼接(concat)起来就得到了新的结果。

图多头注意力的多个输出

直观上可以想到，如果设计这样的一个模型，必然也不会只做一次 attention，多次 attention 综合的结果至少能够起到增强模型的作用，也可以类比 CNN 中同时使用多个卷积核的作用，直观上讲，多头的注意力有助于网络捕捉到更丰富的特征 / 信息。

代码

class MultiHeadAttention(nn.Module):

def __init__(self, dim, dim_qk=None, dim_v=None, num_heads=1, dropout=0.):

super(MultiHeadAttention, self).__init__()

dim_qk = dim if dim_qk is None else dim_qk

dim_v = dim if dim_v is None else dim_v

assert dim % num_heads == 0 and dim_v % num_heads == 0 and dim_qk % num_heads == 0, 'dim must be divisible by num_heads'

self.dim = dim

self.dim_qk = dim_qk

self.dim_v = dim_v

self.num_heads = num_heads

self.dropout = nn.Dropout(dropout)

self.w_q = nn.Linear(dim, dim_qk)

self.w_k = nn.Linear(dim, dim_qk)

self.w_v = nn.Linear(dim, dim_v)

def forward(self, q, k, v, mask=None):

# q: [B, len_q, D]

# k: [B, len_kv, D]

# v: [B, len_kv, D]

assert q.ndim == k.ndim == v.ndim == 3, 'input must be 3-dimensional'

len_q, len_k, len_v = q.size(1), k.size(1), v.size(1)

assert q.size(-1) == k.size(-1) == v.size(-1) == self.dim, 'dimension mismatch'

assert len_k == len_v, 'len_k and len_v must be equal'

len_kv = len_v

q = self.w_q(q).view(-1, len_q, self.num_heads, self.dim_qk // self.num_heads)

k = self.w_k(k).view(-1, len_kv, self.num_heads, self.dim_qk // self.num_heads)

v = self.w_v(v).view(-1, len_kv, self.num_heads, self.dim_v // self.num_heads)

# q: [B, len_q, num_heads, dim_qk//num_heads]

# k: [B, len_kv, num_heads, dim_qk//num_heads]

# v: [B, len_kv, num_heads, dim_v//num_heads]

# The following 'dim_(qk)//num_heads' is writen as d_(qk)

q = q.transpose(1, 2)

k = k.transpose(1, 2)

v = v.transpose(1, 2)

# q: [B, num_heads, len_q, d_qk]

# k: [B, num_heads, len_kv, d_qk]

# v: [B, num_heads, len_kv, d_v]

attn = torch.matmul(q, k.transpose(-2, -1)) / (self.dim_qk ** 0.5)

# attn: [B, num_heads, len_q, len_kv]

if mask is not None:

attn = attn.transpose(0, 1).masked_fill(mask, float('-1e20')).transpose(0, 1)

attn = torch.softmax(attn, dim=-1)

attn = self.dropout(attn)

output = torch.matmul(attn, v)

# output: [B, num_heads, len_q, d_v]

output = output.transpose(1, 2)

# output: [B, len_q, num_heads, d_v]

output = output.contiguous().view(-1, len_q, self.dim_v)

# output: [B, len_q, num_heads * d_v] = [B, len_q, dim_v]

return output

Padding Mask

介绍

由于实际处理时，我们使用的是mini-batch，因此我们会在输入序列中填充（padding），也即是<PAD>token。当我们在输入序列的末尾填充这些无意义的值，以使所有序列达到相同的长度时，这些填充部分会对注意力计算产生不必要的干扰。

Padding Mask 的目的是通过将填充部分的权重设置为负无穷（或者使用非常大的负值）来抑制这些填充部分的影响，使得模型在进行注意力计算时能够忽略这些无效的位置。

具体而言，Padding Mask 是一个与输入序列维度相同的二维矩阵，其中填充部分对应的元素值为 -inf（或者一个很大的负数），而其他部分的值为 0。在注意力计算过程中，将这个 Padding Mask 与注意力分数相加，使得填充部分的注意力分数变成-inf，从而在 Softmax 操作后得到零的权重，最终影响被减弱。

图 Padding Mask

图含Mask的attention

代码

def padding_mask(pad_q, pad_k):

"""

:param pad_q: pad label of query (0 is padding, 1 is not padding), [B, len_q]

:param pad_k: pad label of key (0 is padding, 1 is not padding), [B, len_k]

:return: mask tensor, False for not replaced, True for replaced as -inf

e.g. pad_q = tensor([[1, 1, 0]], [1, 0, 1])

padding_mask(pad_q, pad_q) =

tensor([[[False, False, True],

[False, False, True],

[ True, True, True]],

[[False, True, False],

[ True, True, True],

[False, True, False]]])

"""

assert pad_q.ndim == pad_k.ndim == 2, 'pad_q and pad_k must be 2-dimensional'

assert pad_q.size(0) == pad_k.size(0), 'batch size mismatch'

mask = pad_q.bool().unsqueeze(2) * pad_k.bool().unsqueeze(1)

mask = ~mask

# mask: [B, len_q, len_k]

return mask

残差拼接与正态化(Add&Norm)

残差拼接

为了避免梯度消失，我们在上一步得到了经过 self-attention 加权之后输出，也就是Self-Attention Output ，然后把他们加起来做残差连接。

X = X + Self-Attention Output

正态化(Normalization)

Normalization 的作用是把神经网络中隐藏层归一为标准正态分布，也就是独立同分布，以起到加快训练速度，加速收敛的作用。在Transformer原文中使用的是Layer Normalization，其他文章中也有使用Batch Normalization等其他方式的。

Layer Normalization就是对每个xi 自身各维度进行normalization，与批次无关。具体计算步骤如下：

计算沿特征维度的平均值：

计算沿特征维度的方差：

使用均值和方差对输入进行归一化（其中ϵ 用于避免分母为0）：

对归一化输入进行缩放和移位（可以不含此步骤）：

简单来说就是减去均值之后除以方差，与Batch normalization对比图如下：

图两种Normalization

其中的每一列在Transformer中为一个token的表示或者经过变化后的表示，因此需要注意到Batch Normalization甚至是和输入的其他数据彼此之间有关系的。

前馈神经网络(Feedforward Neural Network)

介绍

Feedforward Neural Network就是一个单隐层的神经网络，通常隐层大小会大于输入层，输入层和输出层大小通常相同。相当于将每个位置的Attention结果映射到一个更大维度的特征空间，然后使用激活函数引入非线性进行筛选，最后恢复回原始维度。在Transformer中，Feedforward Neural Network中的激活函数成为了一个主要的能提供非线性变换的单元。

代码

class Feedforward(nn.Module):

def __init__(self, dim, hidden_dim=2048, dropout=0., activate=nn.ReLU()):

super(Feedforward, self).__init__()

self.dim = dim

self.hidden_dim = hidden_dim

self.dropout = nn.Dropout(dropout)

self.fc1 = nn.Linear(dim, hidden_dim)

self.fc2 = nn.Linear(hidden_dim, dim)

self.act = activate

def forward(self, x):

x = self.act(self.fc1(x))

x = self.dropout(x)

x = self.fc2(x)

return x

Masked Multi-Head Attention

Transformer训练与预测的不同（如何并行计算）

需要注意的是，在Transformer里训练的时候，我们会将真实的答案一同放入到Decoder中，利用真实的答案预测出不同位置的下一个内容。

例如，假如进行翻译，真实的要被翻译出的结果是“我喜欢小狗”，那么我们会放入[<SOS>, 我，喜，欢，小，狗]这样的序列A，而期望输出[我，喜，欢，小，狗，<EOS>]这个序列B。在B中，相当于每个位置的输出都是A中同位置信息以及之前信息共同预测出的下一个token。例如“喜”是由<SOS>和“我”这两个token给出的预测。这也就是说，在训练时，中间某个部分预测错误不会导致后面的预测跟着出错，因为出错误的部分不会再作为输入，可能出现[我，爱，欢，小，狗，<EOS>]这样的错误的句子。

而相对的，用模型做预测时，则是循环输入，先输入[<SOS>]，得到输出[“我”]，接着拼接成[<SOS>，“我”]，得到输出[“我”，“爱”]，再把“爱”拼接成[<SOS>，“我”，“爱”]，以此类推，最终输出<EOS>终止，输出结果为[我，爱，小，狗，<EOS>]，虽然与ground-truth不同，但也是个正确的句子。

Masked Self-Attention

正如上面所说，训练时，我们会把整个[<SOS>, 我，喜，欢，小，狗]都输入进去，但我们又希望后面的结果只能看到之前的内容，总不能希望预测“喜”的时候看到输入是“喜”所以直接输出“喜”，这明显不合适，因此需要一种方法让内容只能看到之前的信息。

具体解决方法就是，在进行 self-attention 操作时，首先通过得到 Scaled Scores，接下来非常关键，我们要对 Scaled Scores 进行 Mask，不应该让它们知道之后词的信息。Mask 非常简单，首先生成一个下三角全 0，上三角全为负无穷的矩阵，然后将其与 Scaled Scores 相加即可。

图计算Masked Self attention

代码

其他部分与Attention相同，只需要增加mask即可：

def attn_mask(len):

"""

:param len: length of sequence

:return: mask tensor, False for not replaced, True for replaced as -inf

e.g. attn_mask(3) =

tensor([[[False, True, True],

[False, False, True],

[False, False, False]]])

"""

mask = torch.triu(torch.ones(len, len, dtype=torch.bool), 1)

return mask

Decoder中的Multi-Head Attention

计算方法与Multi-Head Attention相同，但是K ，和V 来自经历过Encoder之后的Input的表示，而Q 来自Decoder自己的Output。如何理解？

假如翻译：

我喜欢小狗。 | I like dogs.

那么在“I”这个token这里，我们希望它变成“likes”token。“I”对自己进行个变换，成了Q ，去询问当前应该与Input的哪个地方更加契合，之后再把这个地方用来更新自己。“I”发现自己和“喜”“欢”这两个token契合度(Q*K )比较高，因此提取出这两个token的信息(V )更新自己，从而增加了“喜欢”的词义。

总结

上述的内容是单层的Encoder和单层的Decoder的介绍，实际使用中会将这些内容进行堆叠很多层，能发现深层次的信息，这也是大语言模型“大”的一个方面。

Encoder代码

Encoder layer

class EncoderLayer(nn.Module):

def __init__(self, dim, dim_qk=None, num_heads=1, dropout=0., pre_norm=False):

super(EncoderLayer, self).__init__()

self.attn = MultiHeadAttention(dim, dim_qk=dim_qk, num_heads=num_heads, dropout=dropout)

self.ffn = Feedforward(dim, dim * 4, dropout)

self.pre_norm = pre_norm

self.norm1 = nn.LayerNorm(dim)

self.norm2 = nn.LayerNorm(dim)

def forward(self, x, mask=None):

if self.pre_norm:

res1 = self.norm1(x)

x = x + self.attn(res1, res1, res1, mask)

res2 = self.norm2(x)

x = x + self.ffn(res2)

else:

x = self.attn(x, x, x, mask) + x

x = self.norm1(x)

x = self.ffn(x) + x

x = self.norm2(x)

return x

Encoder

class Encoder(nn.Module):

def __init__(self, dim, dim_qk=None, num_heads=1, num_layers=1, dropout=0., pre_norm=False):

super(Encoder, self).__init__()

self.layers = nn.ModuleList([EncoderLayer(dim, dim_qk, num_heads, dropout, pre_norm) for _ in range(num_layers)])

def forward(self, x, mask=None):

for layer in self.layers:

x = layer(x, mask)

return x

Decoder代码

Decoder layer

class DecoderLayer(nn.Module):

def __init__(self, dim, dim_qk=None, num_heads=1, dropout=0., pre_norm=False):

super(DecoderLayer, self).__init__()

self.attn1 = MultiHeadAttention(dim, dim_qk=dim_qk, num_heads=num_heads, dropout=dropout)

self.attn2 = MultiHeadAttention(dim, dim_qk=dim_qk, num_heads=num_heads, dropout=dropout)

self.ffn = Feedforward(dim, dim * 4, dropout)

self.pre_norm = pre_norm

self.norm1 = nn.LayerNorm(dim)

self.norm2 = nn.LayerNorm(dim)

self.norm3 = nn.LayerNorm(dim)

def forward(self, x, enc, self_mask=None, pad_mask=None):

if self.pre_norm:

res1 = self.norm1(x)

x = x + self.attn1(res1, res1, res1, self_mask)

res2 = self.norm2(x)

x = x + self.attn2(res2, enc, enc, pad_mask)

res3 = self.norm3(x)

x = x + self.ffn(res3)

else:

x = self.attn1(x, x, x, self_mask) + x

x = self.norm1(x)

x = self.attn2(x, enc, enc, pad_mask) + x

x = self.norm2(x)

x = self.ffn(x) + x

x = self.norm3(x)

return x

Decoder

class Decoder(nn.Module):

def __init__(self, dim, dim_qk=None, num_heads=1, num_layers=1, dropout=0., pre_norm=False):

super(Decoder, self).__init__()

self.layers = nn.ModuleList([DecoderLayer(dim, dim_qk, num_heads, dropout, pre_norm) for _ in range(num_layers)])

def forward(self, x, enc, self_mask=None, pad_mask=None):

for layer in self.layers:

x = layer(x, enc, self_mask, pad_mask)

return x

Transformer代码

class Transformer(nn.Module):

def __init__(self, dim, vocabulary, num_heads=1, num_layers=1, dropout=0., learnable_pos=False, pre_norm=False):

super(Transformer, self).__init__()

self.dim = dim

self.vocabulary = vocabulary

self.num_heads = num_heads

self.num_layers = num_layers

self.dropout = dropout

self.learnable_pos = learnable_pos

self.pre_norm = pre_norm

self.embedding = nn.Embedding(vocabulary, dim)

self.pos_enc = LearnablePositionalEncoding(dim, 100) if learnable_pos else PositionalEncoding(dim, 100)

self.encoder = Encoder(dim, dim // num_heads, num_heads, num_layers, dropout, pre_norm)

self.decoder = Decoder(dim, dim // num_heads, num_heads, num_layers, dropout, pre_norm)

self.linear = nn.Linear(dim, vocabulary)

def forward(self, src, tgt, src_mask=None, tgt_mask=None, pad_mask=None):

src = self.embedding(src)

src = self.pos_enc(src)

src = self.encoder(src, src_mask)

tgt = self.embedding(tgt)

tgt = self.pos_enc(tgt)

tgt = self.decoder(tgt, src, tgt_mask, pad_mask)

output = self.linear(tgt)

return output

def get_mask(self, tgt, src_pad=None):

# Under normal circumstances, tgt_pad will perform mask processing when calculating loss, and it isn't necessarily in decoder

if src_pad is not None:

src_mask = padding_mask(src_pad, src_pad)

else:

src_mask = None

tgt_mask = attn_mask(tgt.size(1))

if src_pad is not None:

pad_mask = padding_mask(torch.zeros_like(tgt), src_pad)

else:

pad_mask = None

# src_mask: [B, len_src, len_src]

# tgt_mask: [len_tgt, len_tgt]

# pad_mask: [B, len_tgt, len_src]

return src_mask, tgt_mask, pad_mask

引用:
[1]. https://wmathor.com/index.php/archives/1438/

[2]. Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[J]. Advances in neural information processing systems, 2017, 30.

[3]. 基于Keras框架实现加入Attention与BiRNN的机器翻译模型 - 知乎

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/凡人多烦事01/article/detail/566609