编码器层 - 代码实现

import torch
import torch.nn as nn
from dm01_input import *
from dm02_position import *
from dm03_components import *
from dm07_sublayer import SubLayerConnection
from dm04_clones import clones


class EncoderLayer(nn.Module):
    """编码器层"""
    def __init__(self, d_model, self_attn, feed_forward, dropout=0.1):
        """
        初始化编码器层
        :param d_model: 词向量维度，例如 512
        :param self_attn: 多头注意力对象
        :param feed_forward: 前馈全连接对象
        :param dropout: 随机失活概率
        """
        super().__init__()
        
        # 1. 保存子层实例
        self.d_model = d_model
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        
        # 2. 克隆 2 个子层连接结构
        self.sublayer = clones(SubLayerConnection(d_model, dropout), 2)
    
    def forward(self, x, mask):
        """
        前向传播
        :param x: 输入张量 (batch, seq_len, d_model)
        :param mask: 掩码张量
        :return: 编码器层输出
        """
        # 1. 子层1：多头注意力 + Add & Norm
        # 使用 lambda 定义子层函数
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        
        # 2. 子层2：前馈全连接 + Add & Norm
        x = self.sublayer[1](x, self.feed_forward)
        
        return x


# clones 函数
def clones(module, N):
    """克隆 N 个模块"""
    return nn.ModuleList([module for _ in range(N)])


# 测试代码
def dm08_test_encoder_layer():
    """测试编码器层"""
    from dm04_attention import MultiHeadAttention
    from dm05_feedforward import FeedForward
    
    d_model = 512
    
    # 1. 准备数据
    x = use_position()  # (2, 4, 512)
    
    # 2. 实例化子层对象
    multi_head_attn = MultiHeadAttention(d_model, 8)
    feed_forward = FeedForward(d_model, 2048, 0.1)
    
    # 3. 创建编码器层对象
    encoder_layer = EncoderLayer(d_model, multi_head_attn, feed_forward)
    
    # 4. 构建掩码张量
    mask = torch.zeros(2, 1, 4, 4)
    
    # 5. 前向传播
    output = encoder_layer(x, mask)
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    
    return output


def use_position():
    """获取位置编码结果"""
    VOCAB_SIZE = 1000
    D_MODEL = 512
    
    my_embedding = Embeddings(D_MODEL, VOCAB_SIZE)
    x = torch.LongTensor([[1, 2, 3, 4], [5, 6, 7, 8]])
    embedded_x = my_embedding(x)
    
    my_position = PositionalEncoding(D_MODEL, dropout=0.1, max_len=60)
    position_x = my_position(embedded_x)
    
    return position_x


if __name__ == "__main__":
    dm08_test_encoder_layer()