输出部分 - 代码测试

import torch
import torch.nn as nn
import torch.nn.functional as F


class Generator(nn.Module):
    """Transformer 输出部分"""
    
    def __init__(self, d_model, vocab_size):
        super(Generator, self).__init__()
        self.linear = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        return F.log_softmax(self.linear(x), dim=-1)


def use_generator():
    """测试输出部分"""
    print("=== 输出部分 - 代码测试 ===")
    
    # 1. 获取解码器的输出结果
    # result = use_decoder()  # 假设已获取解码器输出
    batch_size = 2
    seq_len = 4
    d_model = 512
    vocab_size = 1000
    
    result = torch.randn(batch_size, seq_len, d_model)  # [2, 4, 512]
    print(f"1. 解码器输出形状: {result.shape}")
    
    # 2. 初始化输出生成器（512 → 1000）
    generator = Generator(d_model, vocab_size)
    print(f"\n2. 输出生成器:\n{generator}")
    
    # 3. 通过生成器获取概率分布
    output = generator(result)
    print(f"\n3. 输出形状: {output.shape}")  # [2, 4, 1000]
    
    # 4. 验证输出维度
    print(f"\n4. 验证输出维度:")
    print(f"   batch_size: {output.shape[0]}")  # 2
    print(f"   seq_len: {output.shape[1]}")  # 4
    print(f"   vocab_size: {output.shape[2]}")  # 1000
    
    # 5. 验证概率和为 1
    print(f"\n5. 验证概率和为 1:")
    
    # 5.1 提取第一个样本第一个词的概率
    log_probs = output[0, 0, :]  # [1000]
    print(f"   第一个样本第一个词的对数概率形状: {log_probs.shape}")
    
    # 5.2 将对数概率转换为概率
    probs = torch.exp(log_probs)  # [1000]
    print(f"   概率形状: {probs.shape}")
    
    # 5.3 验证概率和
    prob_sum = probs.sum().item()
    print(f"   概率和: {prob_sum:.4f}")  # 应该接近 1
    
    # 5.4 验证所有样本的概率和
    all_probs = torch.exp(output)  # [2, 4, 1000]
    prob_sums = all_probs.sum(dim=-1)  # [2, 4]
    print(f"\n   所有样本的概率和:")
    print(f"   {prob_sums}")
    
    # 6. 查看概率最高的词
    print(f"\n6. 查看概率最高的词:")
    topv, topi = output[0, 0, :].topk(5)
    print(f"   第一个样本第一个词概率最高的 5 个词:")
    for i in range(5):
        print(f"   词 {topi[i].item()}: {topv[i].item():.4f}")


# 测试
if __name__ == "__main__":
    use_generator()