如何解决tf.math.confusion_matrix()函数的维度不匹配报错

问题场景 #

处理多分类任务时，尝试用tf.math.confusion_matrix()生成混淆矩阵，代码如下：

y_pred = model.predict(x_test)
y_pred = tf.argmax(y_pred, axis=1)
Y_test = tf.argmax(y_test, axis=1)
matrix = tf.math.confusion_matrix(Y_test, y_pred)

张量输出情况 #

运行后查看张量形状：

• Y_test的输出：

tf.Tensor(
[[0 2 0 ... 0 0 0]
 [0 2 0 ... 0 0 0]
 [0 2 0 ... 0 0 0]
 ...
 [0 0 0 ... 0 0 0]
 [0 0 3 ... 0 0 0]
 [0 0 2 ... 0 0 0]], shape=(2124, 279), dtype=int64)

• y_pred的输出：

tf.Tensor(
[[1 2 2 ... 0 0 0]
 [0 2 3 ... 0 0 0]
 [3 2 0 ... 3 1 3]
 ...
 [3 1 0 ... 2 3 2]
 [1 0 3 ... 1 1 2]
 [1 0 2 ... 1 1 2]], shape=(2124, 279), dtype=int64)

取单个样本查看细节：

• Y_test[1]：

tf.Tensor(
[0 2 0 1 0 3 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0], shape=(279,), dtype=int64)

• y_pred[1]：

tf.Tensor(
[0 2 3 1 3 3 2 2 2 3 2 3 2 3 3 2 1 0 0 0 0 3 1 0 2 3 1 2 0 1 0 0 1 0 0 0 0
 2 0 2 1 0 0 0 0 1 0 0 0 3 2 0 0 3 2 0 0 3 3 0 3 0 0 0 0 1 0 2 1 0 2 3 0 3
 3 0 2 3 1 3 2 0 3 0 0 0 0 0 0 0 0 0 0 0 0 1 3 0 0 0 3 3 0 0 0 0 0 3 0 0 1
 0 3 0 3 3 0 1 0 3 0 0 0 0 0 0 3 0 1 0 0 0 0 0 0 0 0 0 3 0 0 3 0 0 0 0 0 0
 3 0 3 3 0 0 0 3 0 0 0 0 0 0 0 0 0 0 3 0 3 0 0 0 0 3 0 3 0 0 0 0 0 0 0 2 0
 0 1 0 0 0 0 2 0 3 3 0 0 0 0 0 0 0 0 0 0 0 0 0 3 0 0 0 3 0 3 3 0 2 3 0 3 3
 3 3 3 0 3 0 3 0 0 3 0 0 0 0 3 3 3 2 0 0 0 0 0 0 0 0 2 3 0 0 3 0 0 0 3 0 2
 0 0 3 0 0 0 0 0 3 1 2 0 3 2 3 0 3 0 0 0], shape=(279,), dtype=int64)

报错信息 #

运行代码时触发错误：

InvalidArgumentError: Dimensions [0,2) of indices[shape=[2124,2,279]] must match dimensions [0,2) of updates[shape=[2124,279]] [Op:ScatterNd]

问题分析 #

Y_test和y_pred都是**(样本数, 类别数)的二维张量，但tf.math.confusion_matrix()要求输入一维张量——每个样本对应单个类别标签。这说明任务实际是多标签分类**（每个样本可属于多个类别），而非常规单标签多分类，直接套用单标签混淆矩阵方法会导致维度不匹配报错。

解决方法 #

针对多标签分类场景，可选择以下两种方案：

方案1：为每个类别单独生成混淆矩阵 #

遍历每个类别，提取该类别下的真实与预测标签，生成对应混淆矩阵：

num_classes = 279
confusion_matrices = []

for class_idx in range(num_classes):
    # 提取当前类别的二值化标签（是否属于该类）
    true_labels = Y_test[:, class_idx]
    pred_labels = y_pred[:, class_idx]
    # 生成该类别的混淆矩阵
    cm = tf.math.confusion_matrix(true_labels, pred_labels)
    confusion_matrices.append(cm)

# 查看第2类的混淆矩阵示例
print(confusion_matrices[2])

方案2：展平所有标签生成全局混淆矩阵 #

将所有样本的类别标签展平为一维数组，计算全局混淆矩阵以统计所有类别的整体分类情况：

# 展平真实标签与预测标签为一维张量
Y_test_flat = tf.reshape(Y_test, [-1])
y_pred_flat = tf.reshape(y_pred, [-1])

# 生成全局混淆矩阵
global_cm = tf.math.confusion_matrix(Y_test_flat, y_pred_flat)
print(global_cm)

补充说明 #

若任务确实是单标签多分类（每个样本仅属于一个类别），则说明tf.argmax()的axis参数选择错误。单标签分类中，模型输出形状为(样本数, 类别数)，需用axis=1取每个样本概率最大的类别索引，得到一维张量。当前张量形状表明真实标签y_test本身是多标签格式，需先确认任务类型再选择对应方法。

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