如何用TensorFlow.js手动创建混淆矩阵并获取预测标签？

提取预测标签 #

cnn.predict(xTest)返回的是包含每个样本类别概率的张量（以你的墨镜检测二分类任务为例，每个样本对应两个类别的概率值）。要得到预测标签，只需取每个样本概率最大的类别索引：

// 获取预测概率张量
const predictions = cnn.predict(xTest);
// 提取每个样本概率最大的类别索引（即预测标签）
const predictedLabels = tf.argMax(predictions, 1).dataSync();
// 将真实标签转为数组（若yTest是张量格式）
const trueLabels = yTest.dataSync();
// 释放张量内存，避免泄漏
predictions.dispose();

计算混淆矩阵指标 #

遍历真实标签和预测标签数组，统计TP、FP、TN、FN四个值（假设类别1为正类“戴墨镜”，类别0为负类“不戴墨镜”）：

let tp = 0, fp = 0, tn = 0, fn = 0;

for (let i = 0; i < trueLabels.length; i++) {
  const trueLbl = trueLabels[i];
  const predLbl = predictedLabels[i];

  if (trueLbl === 1 && predLbl === 1) tp++;
  if (trueLbl === 0 && predLbl === 1) fp++;
  if (trueLbl === 0 && predLbl === 0) tn++;
  if (trueLbl === 1 && predLbl === 0) fn++;
}

// 输出混淆矩阵结果
console.log(`混淆矩阵：
真阳性(TP): ${tp} | 假阳性(FP): ${fp}
真阴性(TN): ${tn} | 假阴性(FN): ${fn}`);

特殊情况处理 #

如果你的真实标签yTest是one-hot编码格式，需要先转为类别索引：

const trueLabels = tf.argMax(yTest, 1).dataSync();
yTest.dispose();

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