如何从微调后的BERT模型获取多分类任务置信度分数

解决方法 #

你当前的代码直接对模型输出的概率矩阵做了argmax操作，仅提取了概率最高的类别索引，丢失了原始概率值，只需调整代码保留模型输出的原始概率，即可提取对应类别的置信度分数，修改后的代码如下：

import numpy as np

sentences = ["i want to open a bank account"]
pred_tokens = map(tokenizer.tokenize, sentences)
pred_tokens = map(lambda tok: ["<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>] + tok + ["[SEP]"], pred_tokens)
pred_token_ids = list(map(tokenizer.convert_tokens_to_ids, pred_tokens))
pred_token_ids = map(
    lambda tids: tids +[0]*(data.max_seq_len-len(tids)),
    pred_token_ids)
pred_token_ids = np.array(list(pred_token_ids))
# 先保留模型输出的所有类别概率，不直接取argmax
pred_probs = model.predict(pred_token_ids)
predictions = pred_probs.argmax(axis=-1)
for text, label in zip(sentences, predictions):
    # 取对应类别的概率，转换为百分比格式，保留两位小数
    confidence = round(pred_probs[0][label] * 100, 2)
    print("intent:", classes[label])
    print(f"confidence score: {confidence}")
    print()

注意：如果你的模型最后输出层没有添加softmax激活，输出的是logits值，需要先对输出结果做softmax转换再计算置信度：

import tensorflow as tf
pred_probs = tf.nn.softmax(model.predict(pred_token_ids)).numpy()

运行后即可得到你期望的输出格式：

intent: open_account
confidence score: 98.30

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