Transformer底层代码分析

fleur2025-04-23 10:57

transformer架构

主要包括编码器和解码器的内部结构

编码器

编码器包括词元嵌入和位置编码，多头自注意力，前馈层，残差连接和层归一化，以下就是下面的相关代码解析。

位置编码和词元嵌入

将每个单词转化为向量位置，因为Transformer取消了RNN的循环方式进行文本输入，所以使用位置编码（Position Encoding）来表示每个词的位置，向量位置的计算是一种位置（pos）和维度的sin或cos的计算，这样pos和pos+k的位置编码就存在一种线性组合的关系，其实也是一种关系的表示。

偶数维度采用sin函数计算，奇数维度采用cos函数计算，d表示整体维度的大小。将单词与相同长度的向量进行相加，并返回计算后的位置，用于后续的输入。

python 复制代码 
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


# 定义位置编码类，继承自nn.Module，作用将词的位置信息转换为位置向量
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super(PositionalEncoding, self).__init__()
        self.d_model = d_model

        # 创建一个常量的矩阵，
        pe = torch.zeros(max_len, d_model)

        # 对位置进行编码
        for pos in range(max_len):
            # 单次遍历两个位置，偶数位置使用sin函数，奇数位置使用cos函数
            for i in range(0, d_model, 2):
                theta = pos / (10000 ** (i / d_model))
                pe[pos, i] = math.sin(theta)
                pe[pos, i + 1] = math.cos(theta)
        # 扩展维度以便于后续操作，在0的位置加一个新的维度
        pe = pe.unsqueeze(0)
        self.register_buffer('pe',pe)

    def forward(self, x):
        # 对输入的x进行缩放，乘以根号d_model，可以使得单词嵌入表示相对大一些
        x = x * math.sqrt(self.d_model)
        # 获取输入x的长度
        seq_len = x.size(1)
        # 将位置编码矩阵和x进行相加
        x = x + Variable(self.pe[:,:seq_len],requires_grad=False).cuda() 
        return x

自注意力层

上一步中将单词与位置向量叠加作为输入，而在自注意力层，将输入的x转换为Q（Query）查询，K（Query）关键字，V（value）值，用于表示每个单词在上下文中的权重关系，也就是每个单词在此位置所需要的关注程度。 计算的大体过程为：将Q和K进行点积操作，除以缩放因子实现稳定，再用softmax将结果处理，避免梯度爆炸和收敛效率差等问题。之后的结果与V进行点积得出自注意力公式。

为了进一步增强自注意力机制增强聚合上下文信息的能力，提出了多头自注意力机制，就是将经过线形变换的Q，K，V与输入的xi从多维度多个侧面做计算，将结果映射到多组的子空间中。

ini 复制代码 
# 多头自注意力机制
class MutiHeadAttention(nn.Module):
    def __init__(self,heads,d_model,drouout = 0.1):
        super(MutiHeadAttention,self).__init__()

        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads
        self.q_linear = nn.Linear(d_model,d_model)
        self.v_linear = nn.Linear(d_model,d_model)
        self.k_linear = nn.Linear(d_model,d_model)
        self.dropout = nn.Dropout(drouout)
        self.out = nn.Linear(d_model,d_model)

    # d_k表示缩放因子，此处为讲述中的注意力公式计算实现
    def attention(q,k,v,d_k,mask=None,dropout=None):
        # 此处为softmax计算之前的点积缩放计算
        scores = torch.matmul(q,k.transpose(-2,-1)) / math.sqrt(d_k)

        # 掩盖掉那些为了填补长度而增加的单元，使其通过softmax计算后为0
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)

        scores = F.softmax(scores,dim=-1)

        if dropout is not None:
            scores = dropout(scores)

        # socres与value的矩阵乘法
        output = torch.matmul(scores,v)
        return output
    def forward(self,q,k,v,mask=None):
        bs = q.size(0)

        # 线性层，将输入的q、k、v分别映射到d_model维度,分成h个头，每个头维度为d_k
        k = self.k_linear(k).view(bs,-1,self.h,self.d_k)
        q = self.q_linear(q).view(bs,-1,self.h,self.d_k)
        v = self.v_linear(v).view(bs,-1,self.h,self.d_k)

        # 转置操作，方便后续计算
        k = k.transpose(1,2)
        q = q.transpose(1,2)
        v = v.transpose(1,2)

        # 多头注意力计算
        scores = self.attention(q,k,v,self.d_k,mask,self.dropout)
        # 连接多个头并输入到最后线性层
        scores = scores.transpose(1,2).contiguous().view(bs,-1,self.d_model)
        output = self.out(scores)
        return output

前馈层

前馈层接收自注意力层的输出作为输入，并通过Relu激活函数的两层全连接网络对输入做非线性变换，变换对模型的结果产生了很重要的结果。前馈层的隐状态的层数要比自注意力的多。

ruby 复制代码 
class FeedForward(nn.Module):
    def __init__(self,d_model,d_ff=2048,drouout=0.1):
        super(FeedForward,self).__init__()

        # 定义第一个线性层（全连接层）
        # 输入维度为 d_model，输出维度为 d_ff
        # 这个线性层的作用是将输入的特征从 d_model 维度映射到 d_ff 维度
        self.linear1 = nn.Linear(d_model,d_ff)
        # 定义第二个线性层（全连接层）
        # 输入维度为 d_ff，输出维度为 d_model
        # 这个线性层的作用是将经过第一个线性层处理后的特征从 d_ff 维度映射回 d_model 维度
        self.linear2 = nn.Linear(d_ff,d_model)
        # 定义一个 Dropout 层，用于防止过拟合
        # dropout 参数指定了在训练过程中随机丢弃神经元的概率
        self.dropout = nn.Dropout(drouout)

    def forward(self,x):
        # 通过第一个线性层进行线性变换
        # 然后使用 ReLU 激活函数引入非线性
        # ReLU 函数将小于 0 的值置为 0，大于 0 的值保持不变
        # 最后通过 Dropout 层随机丢弃一些神经元的输出，防止过拟合
        x = self.dropout(F.relu(self.linear1(x)))
        # 通过第二个线性层进行线性变换
        # 将经过第一个线性层和 Dropout 层处理后的特征从 d_ff 维度映射回 d_model 维度
        x = self.linear2(x)
        return  x

残差连接和层归一化

Transformer的网络结构庞大又复杂，编码器和解码器包含了很多Transformer块，中间的处理也有非线性变换，为了防止上一层的输出到该层输入过程的误解，采用了残差连接，直接用用直接的线性变换。层归一化借助标准差和均值的方式进行归一处理。

python 复制代码 
# 层归一化处理
class NormLayer(nn.Module):
    def __init__(self,d_model,eps=1e-6):
        super(NormLayer,self).__init__()
        # 定义一个线性层，用于计算均值和方差
        self.size = d_model
        
        # 定义一个可训练的参数，用于调整归一化后的特征
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.beta = nn.Parameter(torch.zeros(self.size))
        self.eps = eps
    def forward(self,x):
        norm = self.alpha * (x - x.mean(dim=-1,keepdim=True)) / (x.std(dim=-1,keepdim=True) + self.eps) + self.beta
        return norm

解码层

ini 复制代码 
# 定义解码器层类，继承自nn.Module
class DecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward):
        super(DecoderLayer, self).__init__()
        # 初始化第一个多头注意力机制（自注意力）
        self.multihead_attention1 = nn.MultiheadAttention(d_model, nhead)
        # 初始化第二个多头注意力机制（与编码器交互）
        self.multihead_attention2 = nn.MultiheadAttention(d_model, nhead)
        # 初始化前馈网络
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model)
        )
        # 初始化第一层归一化层
        self.layer_norm1 = nn.LayerNorm(d_model)
        # 初始化第二层归一化层
        self.layer_norm2 = nn.LayerNorm(d_model)
        # 初始化第三层归一化层
        self.layer_norm3 = nn.LayerNorm(d_model)
        # 初始化dropout层
        self.dropout = nn.Dropout(0.1)

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):
        # 自注意力机制
        attn_output1, _ = self.multihead_attention1(tgt, tgt, tgt, attn_mask=tgt_mask)
        # 残差连接并应用dropout
        tgt = tgt + self.dropout(attn_output1)
        # 归一化
        tgt = self.layer_norm1(tgt)

        # 与编码器输出交互的注意力机制
        attn_output2, _ = self.multihead_attention2(tgt, memory, memory, attn_mask=memory_mask)
        # 残差连接并应用dropout
        tgt = tgt + self.dropout(attn_output2)
        # 归一化
        tgt = self.layer_norm2(tgt)

        # 前馈网络
        ff_output = self.feed_forward(tgt)
        # 残差连接并应用dropout
        tgt = tgt + self.dropout(ff_output)
        # 归一化
        tgt = self.layer_norm3(tgt)
        return tgt

学习资料来自于《大规模语言模型 从理论到实践》张奇 桂韬 郑锐 ⻩萱菁

相关推荐
Baihai_IDP5 小时前
GenAI 时代，数据唾手可得，但真正的挑战已经转变...
人工智能·llm·openai
卓越进步6 小时前
MCP Server架构设计详解：一文掌握框架核心
大模型·llm·go·后端开发·mcp
GetcharZp6 小时前
3小时+10G显卡=你的专属AI助手？MiniMind让语言模型训练"平民化"
llm
架构精进之路7 小时前
大模型重复生成内容：根因剖析与优化策略
后端·llm·ai编程
阿里云大数据AI技术7 小时前
DistilQwen2.5-DS3-0324发布：知识蒸馏+快思考=更高效解决推理难题
人工智能·llm
前端加油站7 小时前
人人都要掌握的 Ollama（奥利玛）
前端·llm
剑客的茶馆17 小时前
GPT，Genini, Claude Llama, DeepSeek，Qwen,Grok,选对LLM大模型真的可以事半功倍！
gpt·llm·llama·选择大模型
微学AI21 小时前
详细介绍：MCP（大模型上下文协议）的架构与组件，以及MCP的开发实践
前端·人工智能·深度学习·架构·llm·mcp
加班挖坑1 天前
本地部署graphRAG
llm·openai
热门推荐
01KGG转MP3工具|非KGM文件|解密音频02从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑03我决定放弃搞 Java 了04YOLOv8入门 | 重要性能衡量指标、训练结果评价及分析及影响mAP的因素【发论文关注的指标】05DeepSeek各版本说明与优缺点分析06RAG 实践- Ollama+RagFlow 部署本地知识库07苍穹外卖面试总结08yolov8，yolo11，yolo12 服务器训练到部署全流程 笔记09如何在WPS和Word/Excel中直接使用DeepSeek功能10组基轨迹建模 GBTM的介绍与实现（Stata 或 R）