Transformer 代码剖析10 - TransformerEmbedding （pytorch实现）

一、模块架构全景图

1.1 核心功能定位

TransformerEmbedding 是 Transformer 架构的输入预处理核心模块，承担着将离散符号序列转化为富含语义和位置信息的连续向量表示的关键任务。

特征融合示意图

# 输入维度：(batch_size, seq_len)
input_tensor = torch.LongTensor([[1, 3, 5], [2, 4, 6]])
 
# 输出维度：(batch_size, seq_len, d_model)
output = TransformerEmbedding(...)(input_tensor)

1.2 模块流程图解

构造函数流程图：

前向传播流程图：

二、代码逐行精解

2.1 类定义与初始化逻辑

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, drop_prob, device):
        super().__init__()  # 继承PyTorch模块特性 
        self.tok_emb = TokenEmbedding(vocab_size, d_model)  # 词嵌入矩阵 
        self.pos_emb = PositionalEncoding(d_model, max_len, device)  # 位置编码器 
        self.drop_out = nn.Dropout(p=drop_prob)  # 正则化装置 

参数矩阵维度分析表

2.2 前向传播动力学

def forward(self, x):
    tok_emb = self.tok_emb(x)  # 符号→向量转换 
    pos_emb = self.pos_emb(x)  # 位置特征注入 
    return self.drop_out(tok_emb + pos_emb)  # 特征融合与正则 

张量变换演示

# 输入 (batch_size=2, seq_len=3)
x = tensor([[5, 2, 8], 
           [3, 1, 0]])
 
# TokenEmbedding输出 (d_model=4)
tok_emb = tensor([[[0.2, 0.5,-0.1, 0.7],
                   [1.1,-0.3, 0.9, 0.4],
                   [0.6, 0.8,-0.2, 1.0]],
                  
                  [[0.9, 0.1, 1.2,-0.5],
                   [0.3, 0.7,-0.4, 0.8],
                   [0.0, 0.0, 0.0, 0.0]]])
 
# PositionalEncoding输出 
pos_emb = tensor([[[0.1, 0.3,-0.2, 0.4],
                   [0.5, 0.1, 0.7,-0.3],
                   [0.2, 0.6, 0.1, 0.9]],
                  
                  [[0.1, 0.3,-0.2, 0.4],
                   [0.5, 0.1, 0.7,-0.3],
                   [0.2, 0.6, 0.1, 0.9]]])
 
# 融合后输出 (dropout_rate=0.1)
output = tensor([[[0.33, 0.88,-0.3, 1.1],  # 保留90%特征 
                 [1.6, -0.2, 1.6, 0.1],
                 [0.8, 1.4, -0.1, 1.9]],
                
                [[1.0, 0.4, 1.0, -0.1],
                 [0.8, 0.8, 0.3, 0.5],
                 [0.2, 0.6, 0.1, 0.9]]])

三、核心子模块原理

3.1 TokenEmbedding 实现机制

• 数学表达：$E_{token}=W_{embed}[X]$
• 训练特性：通过反向传播学习语义关联
• 参数量计算：$|V|\times d_{model}$（V为词汇表）

章节跳转： TokenEmbedding 实现机制解析

3.2 PositionalEncoding 位置编码

• 公式实现：
  $P{E}_{(pos,2i)}=\sin (pos/10000^{2i/d_{model}})$
  $P{E}_{(pos,2i+1)}=\cos (pos/10000^{2i/d_{model}})$

• 优势特性：

  • 相对位置敏感
  • 无限序列扩展性
  • 线性可加性

章节跳转： PositionalEncoding 位置编码实现原理解析

四、关键技术解析

4.1 特征融合策略

tok_emb + pos_emb  # 直接相加而非拼接 

选择依据对比表

4.2 Dropout正则化

nn.Dropout(p=0.1)  # 以10%概率随机置零 

激活模式对比实验

4.3 混合编码机制

设计哲学
1. 解耦设计： 语义与位置信息独立编码
2. 正交性保证： $E_{token}\perp E_{position}$
3. 可扩展性： 支持多种位置编码变体

4.4 动态设备感知

class PositionalEncoding:
    def __init__(self, d_model, max_len, device):
        pe = torch.zeros(max_len, d_model)  # 设备敏感创建 
        self.register_buffer('pe', pe.to(device))

章节跳转： PositionalEncoding 位置编码实现原理解析

五、工程实践要点

5.1 设备兼容性配置

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.pos_emb = PositionalEncoding(..., device)

多设备支持策略

1. 在模块初始化时同步设备状态
2. 使用to(device)方法动态迁移
3. 确保所有子模块设备一致性

5.2 长序列处理机制

max_len = 512  # 典型Transformer设置 

长度扩展方案比较

六、性能优化建议

6.1 内存优化方案

# 使用稀疏梯度优化 
self.tok_emb = nn.Embedding(vocab_size, d_model, sparse=True)

6.2 计算图优化

# 启用PyTorch JIT编译 
@torch.jit.script 
def forward(...):
    ...

原项目代码+注释（附）

"""
@author : Hyunwoong
@when : 2019-10-22
@homepage : https://github.com/gusdnd852
"""

from torch import nn

# 从其他模块导入PositionalEncoding和TokenEmbedding类
from models.embedding.positional_encoding import PositionalEncoding
from models.embedding.token_embeddings import TokenEmbedding

# 定义一个名为TransformerEmbedding的类，它继承自nn.Module
class TransformerEmbedding(nn.Module):
    """
    TransformerEmbedding类结合了词嵌入和正弦位置编码。
    位置编码可以为网络提供单词的位置信息。
    """

    def __init__(self, vocab_size, d_model, max_len, drop_prob, device):
        """
        包含位置信息的词嵌入类的构造函数。

        :param vocab_size: 词汇表的大小。
        :param d_model: 模型的维度，即嵌入向量的维度。
        :param max_len: 序列的最大长度。
        :param drop_prob: Dropout层的丢弃概率。
        :param device: 硬件设备设置（CPU或GPU）。
        """
        super(TransformerEmbedding, self).__init__()  # 调用父类nn.Module的构造函数。
        # 初始化词嵌入层
        self.tok_emb = TokenEmbedding(vocab_size, d_model)
        # 初始化位置编码层
        self.pos_emb = PositionalEncoding(d_model, max_len, device)
        # 初始化Dropout层，用于防止过拟合
        self.drop_out = nn.Dropout(p=drop_prob)

    def forward(self, x):
        """
        前向传播方法，用于计算输入x的嵌入表示。
        """
        # 通过词嵌入层得到词嵌入表示
        tok_emb = self.tok_emb(x)
        # 通过位置编码层得到位置编码表示
        # 注意：这里的位置编码实现可能不是直接应用于x，而是返回一个与x长度相同的位置编码矩阵，然后与tok_emb相加。
        # 正确的实现应该是根据x的序列长度从位置编码矩阵中截取相应部分，但这里为了简化说明，我们假设pos_emb(x)能正确处理。
        pos_emb = self.pos_emb(x)
        # 将词嵌入表示和位置编码表示相加，并通过Dropout层
        return self.drop_out(tok_emb + pos_emb)