Transformer 代码剖析4 - 编码器层实现 （pytorch实现）

一、EncoderLayer - 类结构定义

参考：项目代码

class EncoderLayer(nn.Module):
    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):
        super(EncoderLayer, self).__init__()
        self.attention = MultiHeadAttention(d_model=d_model, n_head=n_head)
        self.norm1 = LayerNorm(d_model=d_model)
        self.dropout1 = nn.Dropout(p=drop_prob)
        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
        self.norm2 = LayerNorm(d_model=d_model)
        self.dropout2 = nn.Dropout(p=drop_prob)

初始化流程图

d_model=512
n_head=8

d_model=512
hidden=2048

eps=1e-6

eps=1e-6

EncoderLayer

参数初始化

构建多头注意力

初始化第一归一化层

设置第一Dropout

构建前馈网络

初始化第二归一化层

设置第二Dropout

维度切分：512→8x64

双线性变换层

稳定数值计算

关键参数解析

1. d_model：特征总维度（默认512），需满足d_model % n_head == 0
2. ffn_hidden：前馈网络中间层维度（通常2048），实现维度扩展->收缩的非线性变换
3. n_head：注意力头数（典型值8），通过并行化学习不同表示子空间
4. drop_prob：正则化概率（常用0.1），在残差连接后实施随机深度衰减

二、前向传播过程逐行解析

def forward(self, x, src_mask):
    # 自注意力阶段
    _x = x  # 保留原始输入
    x = self.attention(q=x, k=x, v=x, mask=src_mask)  # 自注意力计算
    x = self.dropout1(x)  # 随机失活
    x = self.norm1(x + _x)  # 残差连接+层归一化
    
    # 前馈网络阶段
    _x = x  # 保存中间状态
    x = self.ffn(x)  # 位置相关变换
    x = self.dropout2(x)  # 二次正则化
    x = self.norm2(x + _x)  # 最终输出整合
    return x

前向传播流程图

编码器层全流程

并行计算

特征重组

自注意力模块

输入特征矩阵

多头注意力机制

随机失活层1

残差连接1

层归一化1

前馈模块

位置前馈网络

随机失活层2

残差连接2

层归一化2

输出特征矩阵

可在本节末尾附录中查看原项目对应该部分的代码

关键技术细节

1. 自注意力机制：

• 实现公式：$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$
• 掩码处理：通过src_mask实现padding屏蔽和因果约束
• 多头投影：将d_model维度切分为n_head个64维子空间

2. 残差连接：

• 数学表达：$x_{out}=Norm(x+Sublayer(x))$
• 梯度传播：$\frac{\partial loss}{\partial x}=\frac{\partial loss}{\partial x_{out}}+\frac{\partial Sublayer}{\partial x}$
• 初始化策略：各子层采用Xavier正态分布初始化

3. 层归一化：

• 计算式：$y=\frac{x-E[x]}{\sqrt{Var[x]+ϵ}}\ast \gamma +\beta$
• 实施位置：区别于其他架构，采用post-norm方式
• 数值稳定：epsilon=1e-6防止除零错误

4. 前馈网络：

• 结构组成：$FFN(x)=max(0,x{W}_1+b_1)W_2+b_2$
• 维度变换：512→2048→512实现特征重组
• 激活函数：ReLU的非饱和特性防止梯度消失

三、工程实践建议

1. 梯度优化策略

示例：混合精度训练配置
scaler = torch.cuda.amp.GradScaler()
with autocast():
    output = model(input)
scaler.scale(loss).backward()
scaler.step(optimizer)

2. 内存优化技巧

• 使用checkpoint技术节省显存：

from torch.utils.checkpoint import checkpoint
x = checkpoint(self.attention, x, x, x, src_mask)

3. 性能分析工具

使用PyTorch Profiler
with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
    model(input)
print(prof.key_averages().table())

四、扩展改进方向

1. 结构变体：

# 示例：使用Pre-LayerNorm
class PreNormEncoder(EncoderLayer):
    def forward(self, x, mask):
        x = x + self.dropout1(self.attention(self.norm1(x)))
        x = x + self.dropout2(self.ffn(self.norm2(x)))
        return x

2. 注意力优化：

• 线性注意力：$Sim(Q,K)=exp(Q(K^{T}W))$
• 局部窗口：将全局计算限制为滑动窗口

3. 自适应机制：

# 动态Dropout比率
self.dropout1 = nn.Dropout(p=calc_dropout_rate(epoch))

该实现方案在保持原始论文精粹的同时，通过模块化设计和工程优化，实现了生产环境下的高性能Transformer层。实际应用中可根据具体任务需求，灵活调整各组件参数，并结合最新的研究成果进行持续改进。

---

原代码（附）

"""
@author : Hyunwoong
@when : 2019-10-24
@homepage : https://github.com/gusdnd852
# 此部分为代码的作者信息、编写时间和主页链接。
"""

# 从torch库中导入nn模块，nn模块包含了构建神经网络所需的所有组件。
from torch import nn

# 从自定义的models.layers模块中导入LayerNorm、MultiHeadAttention和PositionwiseFeedForward类。
from models.layers.layer_norm import LayerNorm
from models.layers.multi_head_attention import MultiHeadAttention
from models.layers.position_wise_feed_forward import PositionwiseFeedForward

# 定义一个EncoderLayer类，它继承自nn.Module，表示编码器的一层。
class EncoderLayer(nn.Module):

    # 初始化函数，接收模型维度、前馈神经网络的隐藏层维度、多头注意力的头数和dropout概率作为参数。
    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):
        # 调用父类的初始化函数。
        super(EncoderLayer, self).__init__()
        
        # 初始化多头注意力机制，传入模型维度和头数。
        self.attention = MultiHeadAttention(d_model=d_model, n_head=n_head)
        
        # 初始化第一层归一化层，传入模型维度。
        self.norm1 = LayerNorm(d_model=d_model)
        
        # 初始化第一层dropout层，传入dropout概率。
        self.dropout1 = nn.Dropout(p=drop_prob)

        # 初始化位置前馈神经网络，传入模型维度、隐藏层维度和dropout概率。
        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
        
        # 初始化第二层归一化层，传入模型维度。
        self.norm2 = LayerNorm(d_model=d_model)
        
        # 初始化第二层dropout层，传入dropout概率。
        self.dropout2 = nn.Dropout(p=drop_prob)

    # 前向传播函数，接收输入x和源掩码src_mask。
    def forward(self, x, src_mask):
        # 1. 计算自注意力
        _x = x  # 保存原始的输入x，用于后续的残差连接。
        
        # 调用多头注意力机制，传入查询、键、值和掩码，这里查询、键、值都是x。
        x = self.attention(q=x, k=x, v=x, mask=src_mask)
        
        # 2. 加法和归一化
        # 对注意力输出应用dropout。
        x = self.dropout1(x)
        
        # 对dropout后的输出和原始输入进行残差连接，然后进行归一化。
        x = self.norm1(x + _x)
        
        # 3. 位置前馈神经网络
        _x = x  # 保存当前的x，用于后续的残差连接。
        
        # 调用位置前馈神经网络，传入x。
        x = self.ffn(x)
      
        # 4. 加法和归一化
        # 对前馈神经网络的输出应用dropout。
        x = self.dropout2(x)
        
        # 对dropout后的输出和之前的x进行残差连接，然后进行归一化。
        x = self.norm2(x + _x)
        
        # 返回编码器的输出。
        return x