Transformer详解附代码实现

深度学习麋了鹿

345人浏览 · 2026-03-18 14:56:38

深度学习麋了鹿 · 2026-03-18 14:56:38 发布

Transformer

Transformer由论文《Attention is All You Need》提出，接下来我会用通俗的方法解释整体结构流程并且各个模块的pytorch代码实现。

整体结构

Transformer 由 Encoder 和 Decoder 两个部分组成，Encoder 和 Decoder 都包含 6 个 block。

1. Embedding

无论源文本嵌入还是目标文本嵌入，都是为了将文本中词汇的数字表示转变为向量的表示，希望在这样的高维空间捕捉词汇间的关系。其中Embedding方式有很多，我们可以使用Word2Vec等方式得到，也可以使用一个可学习的Embedding进行表示。

代码实现

class Embedding(nn.Module):
    def __init__(self,d_model,vocab) -> None:
        '''
            d_model：表示词嵌入的维度
            vocab：表示词表的大小
        '''
        super().__init__()
        self.d_model = d_model
        self.vocab = vocab
        self.embed = nn.Embedding(vocab,d_model)

    def forward(self,x):
        # x：代表输入进模型的文本通过词汇映射后的数字张量
        return self.embed(x)*math.sqrt(self.d_model)

其中我们会发现embeding之后乘以了 $\sqrt{d_{model}}$ ，为什么？

这个其实主要是为了平衡后面位置编码向量，因为我们初始化embedding的时候服从正态分布 $N(0,\frac{1}{d_{model}})$ ，方差会随着 $d_{model}$ 变大，导致方差变小，取值返回就会比较小，而位置编码使用的是正余弦函数，范围为[-1,1]，相加之后会导致词向量被淹没，语义丢失。

2. PositionalEncoding

在Transformer的编码器结构中，并没有针对词汇位置信息的处理，因此需要在Embedding层后加入位置编码器，将词汇位置不同可能会产生不同语义的信息加入到词嵌入张量中，以弥补位置信息的缺失

公式

其中：

pose：表示单词在句子中的绝对位置

$d_{model}$ ：表示向量嵌入维度，与词向量维度一致

$2i$ ：表示偶数维度

$2i+1$ ：表示奇数维度

优点：

外推性：当我们训练集里面的最大长度为max_len，当来了一个更长的句子，我们可以根据公式直接计算出超出max_len的位置编码
模型可以计算出两个单词位置的相对距离即pos+k可以用pos的位置计算得到，根据正余弦函数的准则： $sin(A+B)=sin(A)cos(B)+cos(A)sin(B)$ 、 $cos(A+B)=cos(A)cos(B)-sin(A)sin(B)$

代码实现

from torch import nn
import torch

class PositionalEncoding(nn.Module):
    def __init__(self,d_model,dropout,max_length) -> None:
        super().__init__()
        '''
            d_model：词嵌入维度
            dropout：置零比例
            max_length：每个句子的最大长度
        '''
        self.dropout = nn.Dropout(dropout)

        # 初始化位置编码矩阵pe，形状为：[max_length, d_model]
        pe = torch.zeros(max_length,d_model)

        # 初始化一个绝对位置矩阵，词汇的绝对位置就是通过它的索引去表示，形状为：[max_length, 1]
        position = torch.arange(0,max_length).unsqueeze(1)

        # 初始化完成绝对位置矩阵，接下来需要考虑如何将绝对位置加入到位置编码矩阵
        # 其中最简单的方式就是我们需要将绝对位置矩阵从[max_length, 1] 变换成[max_length, d_model]形状，然后覆盖位置编码矩阵
        # 做这种变换需要一个1×d_model的矩阵，且这个矩阵能够将位置编码缩放为足够小的数字，方便收敛
        # 初始化一个连续的1×d_model/2的矩阵，然后分别将位置编码矩阵的奇数位置取cos，偶数位置取sin
        div_term = torch.exp(torch.arange(0,d_model,2)*(-(torch.log(torch.Tensor([10000]))/d_model)))
        pe[:,0::2] = torch.sin(position*div_term)
        pe[:,1::2] = torch.cos(position*div_term)

        # 此时得到的pe位置编码矩阵是一个二维矩阵[max_length,d_model]，embedding的输出为一个三维矩阵[batch_size,vec_len,d_model]
        # 因此我们需要在第0维度添加一个维度
        pe=pe.unsqueeze(0)

        # 将位置编码矩阵注册到模型中，不需要跟随模型的优化而变化
        self.register_buffer("pe",pe)

    def forward(self,x):
        # 初始化的max_length是为了能够初始化足够大的位置，在使用过程中需要和文本的输入进行适配，所以需要根据输入进行切分
        emb_pe = self.pe[:,:x.shape[1],:]
        self.pe.requires_grad=False
        return self.dropout(x + emb_pe)

3. Selft-Attention（自注意力机制）

结构

Self-Attention 的结构中我们需要Q(查询),K(键值),V(值)三个矩阵，其中Q、K、V是通过对数据进行线性变换得到。

公式

从公式中我们可以看出，Attention的计算过程就是：Q和K的转置进行内积计算，为了防止内积过大，因此除以 $\sqrt{d_{k}}$ 进行归一化，再经过softmax计算注意力权重，再与V相乘得到。

流程：

input-1的查询向量为[1, 0, 2]，分别乘上input-1、input-2、input-3的键向量，获得三个score为2，4，4。
然后对这三个score取softmax，获得了input-1、input-2、input-3各自的重要程度。
然后将这个重要程度乘上input-1、input-2、input-3的值向量，求和。
获得input-1的输出。

代码实现

import torch

def attention(query,key,value,mask=None,dropout=None):
    # 一般query的最后一个维度是词嵌入的维度
    d_k = query.shape[-1]

    # query和key的转置进行相乘，在除以d_k进行归一化操作
    scores = torch.matmul(query,key.transpose(-2,-1))/torch.sqrt(torch.Tensor([d_k]))

    # 如果mask不为空，mask等于0的位置对应的scores矩阵位置的值改为-1e9，变成很小的值
    if mask is not None:
        scores.masked_fill(mask==0,-1e9)

    # 对scores矩阵的最后一个维度进行softmax操作,得到最终的注意力张量
    p_att = torch.softmax(scores,dim=-1)

    # 如果dropout不为空，则使用dropout操作
    if dropout is not None:
        p_att = dropout(p_att)

    # 将p_att与value相乘，得到最终的query注意力表示，同时返回注意力张量
    return torch.matmul(p_att,value),p_att

4. Encoder结构

Encoder block 结构，可以看到是由 Multi-Head Attention, Add & Norm, Feed Forward模块组成的

4.1 Add & Norm

结构中分为两部分残差和层归一化， X表示Multi-Head Attention或者Feed Forward的输入，其中残差再resnet中提出，可以解决深层网络梯度消失的问题，Layer Normalization 会将每一层神经元的输入都转成均值方差都一样的，这样可以加快收敛。

代码实现

from torch import nn
import torch

class AddNorm(nn.Module):
    def __init__(self,d_model) -> None:
        super().__init__()
        '''
            d_model：词嵌入维度
            dropout：置零比例
            max_length：每个句子的最大长度
        '''
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self,x,model_out):
        return self.layer_norm(x+model_out)

4.2 MultiHeadAttention

公式

多头注意力机制：核心就是将 $d_{model}$ 分割成N份，然后各自去进行self-attention的计算，为了更加直观就是假设我们输入的QKV的向量形状为[batch_size，seq_len， d_model]，然后我们将d_model分成N份，即每一份的形状为[batch_size，seq_len， N，d_model/N]，然后将N和seq_len进行维度变换得到[batch_size， N，seq_len，d_model/N]，然后计算self-attention，计算完成后再将N和seq_len维度还原，再将N和d_model/N拼接还原为d_model。

优点：能够让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能出现的偏差，让词义拥有来自更多元的表达。

代码实现

import torch
from torch import nn
import copy

# 多头注意力机制
# 实现克隆函数，因为在多头注意力机制下需要用到多个结构相同的线性层
# 需要使用clone函数将他们一同初始化到一个网络层列表中
def clones(module,N):
    # module: 代表要克隆的目标网络层
    # N: 克隆几个
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

# 实现多头注意力机制的类
class MultiHeadAttention(nn.Module):
    def __init__(self,head,embedding_dim,dropout=0.1) -> None:
        super().__init__()
        '''
            head: 代表几个头的参数
            embedding: 代表词嵌入的维度
            dropout: 置零的比率
        '''
        # 断言多头的数量能够整除词嵌入的维度embedding_dim
        assert embedding_dim%head ==0

        # 获得每个头词向量的维度
        self.d_k = embedding_dim//head

        self.head = head
        self.embedding_dim = embedding_dim

        # 获得线性层，要获取4个，分别为Q、K、V以及最终的输出线性层
        self.linears = clones(nn.Linear(embedding_dim,embedding_dim),4)

        # 初始化注意力张量
        self.att = None

        # 初始化dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self,query,key,value,mask=None):
        # query,key,value是注意力机制的三个输入张量，mask为掩码张量
        if mask is not None:
            # 使用unsqueeze将掩码张量进行维度扩充，代表多头中的第n个头
            mask = mask.unsqueeze(1)
        
        # 得到batch_size
        batch_size = query.shape[0]

        # 首先使用zip将网络层和输入数据连接在一起，模型的输出利用view和transpose进行维度和形状的变换输出
        # 将句子长度的维度和头数维度进行交换目的更好的学习和表达句子长度和词嵌入维度之间的关系
        query,key,value = [
            model(x).view(batch_size,-1,self.head,self.d_k).transpose(1,2) for model,x in zip(self.linears,(query,key,value))
        ]

        # 将每个头的输出传入到注意力层
        x,self.attn = attention(query,key,value,mask=mask,dropout=self.dropout)

        # 得到每个头的计算结果是4维张量，需要进行形状的转换
        # 前面已经将1,2两个维度进行转置，需要重新转置回来
        # 经历的transpose方法后，必须使用contiguous，不然无法使用view()方法
        x = x.transpose(1,2).contiguous().view(batch_size,-1,self.head*self.d_k)

        # 最后将x输入线性层列表中的最后一个线性层中，得到最终的多头注意力结构输出
        return self.linears[-1](x)

4.3 FeedForward

Feed Forward 层比较简单，是一个两层的全连接层再加一层 Relu激活函数。

公式

代码实现

from torch import nn
import torch

class FeedForward(nn.Module):
    def __init__(self,d_model,d_hidden,dropout=0.1) -> None:
        '''
            d_model: 词嵌入维度
            d_hidden: 隐藏层维度
            dropout: 置零
        '''
        super().__init__()
        self.w1 = nn.Linear(d_model,d_hidden)
        self.w2 = nn.Linear(d_hidden,d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self,x):
        # 首先x先送入第一个全连接，经过relu激活、dropout
        # 最后送入第二个全连接
        return self.w2(self.dropout(torch.relu(self.w1(x))))

5. Decoder结构

5.1 Masked MultiHeadAttention

Decoder中的第一个多头注意力层是有一个Masked操作，主要是在预测下一个词的时候我们只能看到前面的信息，无法看到后面的信息，需要把后面的信息进行掩盖。预测过程如下：

首先根据输入 "<begin>" 预测出第一个单词为 "我"，然后根据输入 "<begin> 我" 预测下一个单词 "爱"，依次知道遇到<end>结束符预测结束。

训练阶段为什么要使用masked？

答：因为在我们训练过程中，在计算loss时，是用当前decoder输入所有单词对应位置的输出 $y_1,y_2,y_3...y_n$ 与真实的翻译结果ground truth去分别算cross entropy loss，然后把t个loss加起来的，如果我们不加mask，那么就会出现在输出 $y_1$ 的时候使用了 $x_1,x_2,x_3,...x_n$ 的信息，其中包含了我们要预测的 $x_2$ 信息，这个在我们推理过程中是不可能给到的，所以我们需要把当前预测位置后面的未来信息进行masked。

mask矩阵如何构建？

答：mask矩阵就是一个下三角矩阵，这样我们会发现对于每一个token都只能观察到当前token以及之前的信息，未来的信息都会被遮盖，mask在计算注意力分数的softmax之前进行使用，将mask==0的位置赋一个 $-e^{9}$ 很小的值，这样在softmax计算时未来的token注意力分数无限小趋近0。

代码实现

def subsquent_mask(size):
    '''
        生成向后遮掩的掩码张量，参数size是掩码张量最后两个维度的大小，
        最后两个维度形成一个方阵
    '''
    # 定义掩码的张量
    att_shape = (1,size,size)
    # 生成上三角矩阵，其中参数"1"表示上三角的对角线向右上移动一个单位，"0"表示不动，"-1"表示向左下移动一个单位
    subsquent_mask = torch.triu(torch.ones(att_shape,dtype=torch.uint8),1)

    # 使用1-上三角矩阵，得到一个下三角矩阵
    return 1-subsquent_mask

5.2 MultiHeadAttention

Decoder block 第二个 Multi-Head Attention 变化不大，主要的区别在于其中 Self-Attention 的 K, V矩阵不是使用上一个 Decoder block 的输出计算的，而是使用 Encoder 的输出计算得到 K、V，根据上一个 Decoder block 的输出计算 Q (如果是第一个 Decoder block 则使用输入矩阵 X 进行计算)。这样做的好处是在 Decoder 的时候，每一位单词都可以利用到 Encoder 所有单词的信息 (这些信息无需 Mask)。

代码实现

同4.2一致

5.3 Softmax输出

Softmax 根据输出矩阵的每一行预测下一个单词：

代码实现

class Generator(nn.Module):
    def __init__(self,d_model,vocab_size) -> None:
        '''
            d_model:词嵌入维度
            vocab_size:词表大小
        '''
        super().__init__()

        self.project = nn.Linear(d_model,vocab_size)

    def forward(self,x):
        return torch.softmax(self.project(x),dim=-1)

AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

《GodoOS全开源与零配、操作简单直观：适合所有用户看看打工人的“哆啦A梦”》

AtomGit开源社区

空间组学下一代机器学习与深度学习

AtomGit开源社区

C++ CMake构建系统最佳实践：模块化与跨平台构建

本文总结了C++项目中使用CMake构建系统的最佳实践。基于AIDC项目的实践经验，文章详细介绍了模块化项目结构设计、核心CMake配置方法、跨平台构建技巧等内容。重点包括：1) 采用分层模块化结构管理项目；2) 通过目标属性和现代CMake特性优化构建过程；3) 使用FetchContent等工具实现依赖管理；4) 支持多种编译器和平台的构建配置。文章还提供了完整的CMakeLists.txt示