输出部分 - 代码实现

1 课程概览

本课讲解 Transformer 输出部分的代码实现。输出部分由一个线性层和一个 LogSoftmax 组成。线性层将解码器的输出从 d_model（512）维转换为 vocab_size（1000）维，LogSoftmax 将分数转换为对数概率分布。

2 核心概念与定义

Generator 类：Transformer 的输出部分。
线性层（Linear）：将 d_model（512）维转换为 vocab_size（1000）维。
LogSoftmax：将分数转换为对数概率分布。
d_model：词向量维度（512）。
vocab_size：词汇表大小（1000）。

3 模型与算法详解

输出部分结构

解码器输出 [2, 4, 512] → Linear → [2, 4, 1000] → LogSoftmax → 概率分布 [2, 4, 1000]

形状变化

步骤	输入形状	输出形状	说明
Linear	[2, 4, 512]	[2, 4, 1000]	线性变换
LogSoftmax	[2, 4, 1000]	[2, 4, 1000]	转概率

参数说明

参数	值	说明
d_model	512	词向量维度
vocab_size	1000	词汇表大小

输出含义

[2, 4, 1000] 表示 2 个句子，每个句子 4 个单词，每个单词是 1000 个词中的概率分布。

[2, 4, 1000]
  ↓
2：2 个句子
4：每个句子 4 个单词
1000：每个单词是 1000 个词中的概率

4 数学原理与推导

线性层

$$\text{output} = \text{Linear}(\text{x})$$

其中：

$\text{x}$ 是解码器输出 [batch, seq_len, d_model]
$\text{Linear}$ 将 d_model 维映射到 vocab_size 维

LogSoftmax

$$\text{output} = \text{LogSoftmax}(\text{output})$$

其中：

$\text{dim}=1$ 表示在第 1 维（vocab_size 维）上做 softmax

5 代码示例

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):
        """
        初始化函数
        
        Args:
            d_model: 词向量维度（512）
            vocab_size: 词汇表大小（1000）
        """
        super(Generator, self).__init__()
        
        # 1. 全连接层：512 → 1000
        self.linear = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        """
        前向传播
        
        Args:
            x: 解码器输出 [batch, seq_len, d_model]
        
        Returns:
            output: 概率分布 [batch, seq_len, vocab_size]
        """
        # 1. 线性层：[2, 4, 512] → [2, 4, 1000]
        x = self.linear(x)
        
        # 2. LogSoftmax：[2, 4, 1000] → [2, 4, 1000]
        # dim=-1 表示在最后一维（vocab_size 维）上做 softmax
        return F.log_softmax(x, dim=-1)


# 测试
if __name__ == "__main__":
    print("=== 输出部分 - 代码实现 ===")
    
    d_model = 512
    vocab_size = 1000
    
    # 创建输出部分
    generator = Generator(d_model, vocab_size)
    print(f"输出部分:\n{generator}")
    
    # 创建输入（解码器输出）
    batch_size = 2
    seq_len = 4
    x = torch.randn(batch_size, seq_len, d_model)  # [2, 4, 512]
    print(f"\n输入形状: {x.shape}")
    
    # 前向传播
    output = generator(x)
    
    print(f"输出形状: {output.shape}")  # [2, 4, 1000]
    
    # 验证概率和为 1
    # exp(log_softmax) = softmax
    probs = torch.exp(output)
    print(f"\n第一个样本第一个词的概率和: {probs[0, 0, :].sum().item():.4f}")
    
    # 查看概率最高的词
    topv, topi = output[0, 0, :].topk(5)
    print(f"\n第一个样本第一个词概率最高的 5 个词:")
    for i in range(5):
        print(f"  词 {topi[i].item()}: {topv[i].item():.4f}")

代码说明

代码	说明
`nn.Linear(d_model, vocab_size)`	全连接层（512 → 1000）
`F.log_softmax(x, dim=-1)`	LogSoftmax
`torch.exp(output)`	将对数概率转换为概率
`output.topk(5)`	取概率最高的 5 个词