如何将HuggingFace Transformers BERT Keras模型输出logit转概率

无需改动已训练模型的快速实现 #

你当前模型输出原始logit的核心原因，是训练时损失函数设置了from_logits=True——这一参数会把softmax的计算逻辑融合进损失层，并未在网络输出端接入激活函数，因此无需重新训练，直接在推理阶段对logit做softmax转换即可。

首先引入TensorFlow的神经网络模块：
import tensorflow as tf

执行预测并转换为概率的完整代码：

# 调用现有模型拿到原始logit
raw_logits = model.predict(x={'input_ids': x['input_ids']})
# 沿样本维度做softmax，输出范围0-1
prob_output = tf.nn.softmax(raw_logits['issue'], axis=1).numpy()

• 参数axis=1对应多分类场景下的类别维度，保证单样本内所有类别的概率和为1
• 输出的prob_output维度与训练标签完全对齐，每个取值都落在0-1区间

如果需要拿到最终预测标签，配合之前的标签编码器直接解码即可：

import numpy as np
predicted_label = le.inverse_transform(np.argmax(prob_output, axis=1))

构建专用推理模型，predict直接输出概率 #

如果不想每次预测后手动计算softmax，可以直接在现有模型的输出端追加softmax层，整个过程共享原有训练权重，不需要重新训练：

# 在原模型的logit输出后接入softmax激活
softmax_output = tf.keras.layers.Softmax(name='issue_proba')(model.output['issue'])
# 构建推理专用的新模型
infer_model = tf.keras.Model(inputs=model.input, outputs=softmax_output)
# 后续调用推理模型直接拿到0-1范围的概率
prob_output = infer_model.predict(x={'input_ids': x['input_ids']})

注意：该推理模型与原训练模型共享全部层参数，仅在推理链路新增激活计算，不会改动原有训练得到的权重

• 不要在训练阶段把softmax接入网络输出，同时保留损失函数的from_logits=True配置，这种写法会重复计算softmax，导致训练收敛异常，你当前的训练配置是合理的
• 不要混用softmax和sigmoid：你当前的任务是单标签多分类（单样本对应一个固定类别），使用softmax是正确选择；如果是多标签任务（单样本可对应多个类别）才需要使用sigmoid将每个类别的logit独立映射到0-1区间
• 当前代码未将attention_mask传入模型，仅使用input_ids训练时，模型会将padding填充的token纳入有效计算范围，不影响概率转换逻辑，但补全attention_mask输入可以有效提升模型精度

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