PyTorch Softmax输出多值而非预期真假概率问题排查

核心问题本质 #

你遇到的情况是因为模型输出维度不符合二分类任务预期，或者softmax的作用维度设置错误，导致输出了序列中每个token的概率分布，而非整段文本的fake/real二分类结果。以下是具体排查点：

1. 模型输出维度不对 #

• 二分类任务要求模型最后一层输出2维张量（对应fake/real两类），但GPT类模型默认输出是(batch_size, seq_len, vocab_size)这类带序列维度的张量，直接做softmax会得到每个token的多分类结果，而非整句的二分类概率。
• 自定义微调GPT2时，必须在模型顶部加二分类头（比如nn.Linear(model.config.hidden_size, 2)），否则模型输出还是语言模型的下一个token预测结果，不是二分类特征。

2. Softmax维度设置错误 #

• 要得到整段文本的二分类概率，得先对模型输出做维度压缩：比如取最后一个token的输出output[:, -1, :]，或者做全局平均池化，把序列维度压成单个样本的特征，再对这个2维特征做softmax（dim=-1）。
• 错误操作：直接对(batch, seq_len, 2)的张量做softmax，会得到每个位置的fake/real概率，不是整段文本的结果。

3. ONNX推理时的张量问题 #

• 导出ONNX模型时要指定固定的batch_size和seq_len，否则推理时可能输出动态维度的张量，导致后续处理混乱。
• 打印推理后的输出形状print(output.shape)，确认是不是(batch_size, 2)，如果是三维张量，说明没做维度压缩。

4. 检测器模型的使用方式错了 #

• OpenAI的检测器是针对整段文本的二分类模型，不能直接拿它的原始输出做softmax，得按照模型的预期流程处理：先获取文本的特征表示，再经过分类头得到二分类结果，最后做softmax。

快速验证步骤 #

• 打印模型输出的张量形状，确认是否有多余的序列维度。
• 提取最后一个token的输出，再做softmax：prob = torch.softmax(output[:, -1, :], dim=-1)，看是否得到两个加和为1的数值。
• 检查微调代码，确认是否添加了正确的二分类头。

内容的提问来源于stack exchange，提问作者Jesper Ezra