子层结构搭建 - 代码测试

def sublayer(x):
    # 创建多头注意力对象
    multi_attention = MultiHeadAttention(512, 8)
    # 计算注意力（自注意力机制 Q=K=V=x）
    return multi_attention(x, x, x)

import torch
from dm01_input import Embeddings, PositionalEncoding
from dm04_attention import MultiHeadAttention
from dm07_sublayer import SubLayerConnection


def use_sublayer():
    """测试子层连接结构"""
    # 1. 准备输入数据（词嵌入+位置编码）
    x = use_position()  # (2, 4, 512)
    
    # 2. 实例化子层连接结构
    sublayer_connection = SubLayerConnection(512)
    
    # 3. 定义子层函数（用多头注意力作为示例）
    # 采用函数嵌套（闭包）写法
    def sublayer(x):
        # 3.1 创建多头注意力对象
        multi_attention = MultiHeadAttention(512, 8)
        # 3.2 计算注意力（自注意力机制 Q=K=V=x）
        return multi_attention(x, x, x)
    
    # 4. 通过子层连接处理输入
    result = sublayer_connection(x, sublayer)
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {result.shape}")
    
    return result


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__":
    result = use_sublayer()
    print(f"最终结果形状: {result.shape}")

def sublayer(x):
    multi_attention = MultiHeadAttention(512, 8)
    return multi_attention(x, x, x)