如何从Hugging Face文本分类pipeline获取logits或原始模型分数？

你当前用默认sentiment-analysis pipeline只能拿到单类别或所有类别的softmax后概率，要获取原始logits（softmax前的输出）或完整概率分布，可通过以下两种方式实现：

方法一：直接使用Tokenizer与Model（推荐） #

跳过封装好的pipeline，直接调用模型和tokenizer，能完整获取模型输出的logits，同时也能自行计算完整概率分布：

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# 加载目标模型与对应的tokenizer
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 处理输入文本
text = "I love this movie!"
inputs = tokenizer(text, return_tensors="pt")

# 禁用梯度计算，提升推理效率
with torch.no_grad():
    outputs = model(**inputs)

# 获取softmax前的原始logits
raw_logits = outputs.logits
print("原始Logits:", raw_logits)

# 计算所有类别的概率分布（手动执行softmax）
prob_distribution = torch.nn.functional.softmax(raw_logits, dim=-1)
print("完整类别概率分布:", prob_distribution)

# 映射标签ID到标签名称，输出每个类别的概率
id2label = model.config.id2label
for idx, prob in enumerate(prob_distribution[0]):
    print(f"{id2label[idx]}: {prob.item():.4f}")

方法二：修改Pipeline配置返回完整结果 #

如果想继续使用pipeline，可通过return_all_scores=True获取所有类别的softmax概率，若需要logits则需额外调用模型：

from transformers import pipeline
import torch

# 创建pipeline时指定return_all_scores=True
clf = pipeline("sentiment-analysis", return_all_scores=True)
result = clf("I love this movie!")
print("所有类别softmax概率:", result)

# 直接调用pipeline内置的模型获取logits
text = "I love this movie!"
inputs = clf.tokenizer(text, return_tensors="pt")
with torch.no_grad():
    outputs = clf.model(**inputs)
print("原始Logits:", outputs.logits)

关键说明 #

• return_all_scores=True仅返回所有类别的softmax后概率，而非原始logits，这就是你之前设置后仍得不到logits的原因。
• 原始logits是模型输出的未经过激活函数（softmax）的原始值，适合用于分数校准、温度缩放这类需要原始输出的操作；完整概率分布则可直接用于计算置信区间。

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