DistilBert二分类遇InvalidArgumentError及指标计算问题求助

基于DistilBert的二分类项目（垃圾短信/IMDB数据集均复现问题），计算召回率、精确率和AUC时触发InvalidArgumentError，错误信息如下：

InvalidArgumentError: predictions must be >= 0
Condition x >= y did not hold.
First 3 elements of x:
[-9.  9. -9.]
First 1 elements of y:
[0.]

模型使用BinaryCrossentropy(from_logits=True)损失函数与Adam优化器，预测输出为[-9., 9.]这类负值（本应是类别概率），在计算召回率步骤触发错误。数据集标签已做One-hot编码（0代表正常短信，1代表垃圾短信），完整代码如下：

import pandas as pd
import tensorflow as tf
import transformers
from transformers import DistilBertTokenizer
from transformers import TFAutoModelForSequenceClassification
pd.set_option('display.max_colwidth', None)
MODEL_NAME = 'distilbert-base-uncased'
BATCH_SIZE = 8
N_EPOCHS = 3

train = pd.read_csv("train_set.csv", error_bad_lines=False)
test = pd.read_csv("test_set.csv", error_bad_lines=False)

X_train = train.text
X_test = test.text
y_train = train.label
y_test = test.label

#One-hot encoding of labels
y_train_encoded = tf.one_hot(y_train.values, 2)
y_test_encoded = tf.one_hot(y_test.values, 2)

tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)

train_encodings = tokenizer(list(X_train.values),
                        truncation=True, 
                        padding=True)
test_encodings = tokenizer(list(X_test.values),
                       truncation=True, 
                       padding=True)

train_dataset = tf.data.Dataset.from_tensor_slices((dict(train_encodings),list(y_train_encoded)))

test_dataset = tf.data.Dataset.from_tensor_slices((dict(test_encodings),list(y_test_encoded)))
test_dataset2 = test_dataset.shuffle(buffer_size=1024).take(1000).batch(16)

model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

optimizerr = tf.keras.optimizers.Adam(learning_rate=5e-5)

losss = tf.keras.losses.BinaryCrossentropy(from_logits=True)

model.compile(optimizer=optimizerr,
          loss=losss,
          metrics=['accuracy'])

print("Evaluate Base model on test data")
results = model.evaluate(test_dataset2)
print("test loss, test acc:", results)

model.fit(train_dataset.shuffle(len(X_train)).batch(BATCH_SIZE),
      epochs=N_EPOCHS,
      batch_size=BATCH_SIZE)

predictions = model.predict(test_dataset2)
# predictions = TFSequenceClassifierOutput(loss=None, logits=array([[-8.96754 ,  8.966875],....[-8.970767,  8.975923]], dtype=float32), hidden_states=None, attentions=None)

import numpy as np
rounded_predictions = np.rint(predictions.logits)
# rounded_predictions = array([[-9.,  9.],....[-9.,  9.]], dtype=float32)

y_test_encoded
#y_test_encoded = <tf.Tensor: shape=(1930, 2), dtype=float32, numpy=array([[1., 0.],....[0., 1.]], dtype=float32)>

m = tf.keras.metrics.Recall()
m.update_state(y_test_encoded, rounded_predictions)
m.result().numpy()
# This is where I getting the above mentioned error.

核心问题分析 #

模型输出的是logits（原始预测值，未经过激活函数转换），不是0-1区间的概率值。直接对logits取整会得到负值，而TensorFlow的召回率等指标要求输入必须是合法的概率或类别标签，因此触发错误。

1. 将Logits转换为类别概率 #

由于使用了BinaryCrossentropy(from_logits=True)，需要通过Softmax激活函数将logits转换为符合要求的类别概率（二分类+One-hot标签场景下，Softmax适配性更好）：

# 替换原预测处理代码
from tensorflow.keras.activations import softmax

# 将logits转换为0-1区间的概率
pred_probs = softmax(predictions.logits, axis=1).numpy()

2. 生成合法输入计算指标 #

召回率、精确率计算

需要将One-hot编码的真实标签转回一维标签，同时从概率中生成预测标签（取概率最大的类别）：

# 真实标签：从One-hot转回一维（0=正常，1=垃圾）
true_labels = np.argmax(y_test_encoded, axis=1)
# 预测标签：取概率最大的类别索引
pred_labels = np.argmax(pred_probs, axis=1)

# 计算召回率
recall_metric = tf.keras.metrics.Recall()
recall_metric.update_state(true_labels, pred_labels)
print(f"召回率: {recall_metric.result().numpy():.4f}")

# 计算精确率
precision_metric = tf.keras.metrics.Precision()
precision_metric.update_state(true_labels, pred_labels)
print(f"精确率: {precision_metric.result().numpy():.4f}")

AUC值计算

AUC需要输入正类（垃圾短信，类别1）的概率值，直接取概率数组的第二列即可：

# 提取类别1的概率值
pos_probs = pred_probs[:, 1]
# 计算AUC
auc_metric = tf.keras.metrics.AUC()
auc_metric.update_state(true_labels, pos_probs)
print(f"AUC值: {auc_metric.result().numpy():.4f}")

3. 可选：编译时直接添加指标 #

如果不想手动计算，可在模型编译阶段直接指定需要的指标，训练/评估时自动输出结果：

model.compile(
    optimizer=optimizerr,
    loss=losss,
    metrics=['accuracy', tf.keras.metrics.Recall(), tf.keras.metrics.Precision(), tf.keras.metrics.AUC()]
)

# 评估时直接获取所有指标结果
results = model.evaluate(test_dataset2)
print("测试集指标结果:", dict(zip(model.metrics_names, results)))

完整修改后的预测与指标计算代码片段 #

predictions = model.predict(test_dataset2)

# 转换logits为类别概率
pred_probs = tf.keras.activations.softmax(predictions.logits, axis=1).numpy()

# 转换标签格式
true_labels = np.argmax(y_test_encoded, axis=1)
pred_labels = np.argmax(pred_probs, axis=1)
pos_probs = pred_probs[:, 1]

# 计算召回率
recall = tf.keras.metrics.Recall()
recall.update_state(true_labels, pred_labels)
print(f"召回率: {recall.result().numpy():.4f}")

# 计算精确率
precision = tf.keras.metrics.Precision()
precision.update_state(true_labels, pred_labels)
print(f"精确率: {precision.result().numpy():.4f}")

# 计算AUC
auc = tf.keras.metrics.AUC()
auc.update_state(true_labels, pos_probs)
print(f"AUC值: {auc.result().numpy():.4f}")

内容的提问来源于stack exchange，提问作者Injarapu Sri Sharanya