如何通过HuggingFace文本分类pipeline获取模型logits？

针对你使用distilbert-base-uncased-finetuned-sst-2-english模型，通过Pipeline执行推理时无法获取logits的问题，有两种可行方案，其中无需自定义Pipeline的方式是最简方案，具体如下：

方案一：直接复用Pipeline加载的模型与Tokenizer（最简） #

Pipeline本质上帮你封装了模型和Tokenizer的加载流程，你可以直接调用其内置的model和tokenizer属性，沿用你之前获取logits的逻辑：

from transformers import pipeline
import torch

# 创建文本分类Pipeline
classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

# 测试文本
texts = ["this is the first sentence", "of my data.", "In fact, thats not true,", "but we are going to assume it", "is"]

# 用Pipeline的Tokenizer处理文本（自动适配模型的padding/truncation规则）
tokenized_test = classifier.tokenizer(
    texts,
    padding=True,
    truncation=True,
    return_tensors="pt"
)

# 获取logits
with torch.no_grad():
    logits = classifier.model(**tokenized_test).logits

# 输出结果
print(logits)

这种方式完全复用Pipeline的配置，省去了手动加载模型和Tokenizer的步骤，同时直接得到原始logits，是最快捷的实现方式。

方案二：自定义Pipeline（适合需长期返回logits的场景） #

如果需要每次调用Pipeline时都自动返回logits（而非单独编写获取逻辑），可以自定义Pipeline类，重写后处理方法来添加logits字段：

from transformers import TextClassificationPipeline
from transformers.pipelines.text_classification import AggregationStrategy

class LogitsTextClassificationPipeline(TextClassificationPipeline):
    def _postprocess(self, model_outputs, function_to_apply=None, return_all_scores=False, aggregation_strategy=AggregationStrategy.NONE):
        # 保留原始logits
        logits = model_outputs["logits"]
        # 调用父类的默认后处理逻辑，得到label和score
        results = super()._postprocess(model_outputs, function_to_apply, return_all_scores, aggregation_strategy)
        # 为每个结果添加logits字段
        for idx, result in enumerate(results):
            result["logits"] = logits[idx].tolist()
        return results

# 创建自定义Pipeline
classifier = LogitsTextClassificationPipeline(model="distilbert-base-uncased-finetuned-sst-2-english")

# 推理并获取包含logits的输出
outputs = classifier(texts)
for item in outputs:
    print(item)

执行后，每个输出项会同时包含label、score和logits三个字段，适合需要将logits作为固定输出的长期使用场景。

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