如何在手动拆分数据集的CNN模型中获取混淆矩阵等评估指标及图表

针对你手动划分数据集且已训练好CNN的场景，以下是添加混淆矩阵、Precision/Recall/F1、ROC曲线及AUC的具体实现：

1. 提取测试集的真实标签与预测结果 #

首先将模型切换到评估模式，遍历测试集获取真实标签、预测类别及预测概率：

import torch
import numpy as np
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score, roc_curve, auc
from sklearn.preprocessing import label_binarize
import matplotlib.pyplot as plt

# 替换为你的模型和测试集加载器
model = your_trained_cnn_model
test_dataloader = your_test_dataloader
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.eval()
true_labels = []
pred_probs = []
pred_labels = []

# 关闭梯度计算，节省资源
with torch.no_grad():
    for inputs, labels in test_dataloader:
        inputs = inputs.to(device)
        outputs = model(inputs)
        
        # 多分类用softmax，二分类替换为torch.sigmoid(outputs)
        prob = torch.softmax(outputs, dim=1)
        pred_label = torch.argmax(prob, dim=1)
        
        # 将张量转numpy并存入列表
        true_labels.extend(labels.cpu().numpy())
        pred_probs.extend(prob.cpu().numpy())
        pred_labels.extend(pred_label.cpu().numpy())

# 转为numpy数组方便后续计算
true_labels = np.array(true_labels)
pred_probs = np.array(pred_probs)
pred_labels = np.array(pred_labels)

2. 计算混淆矩阵、Precision、Recall、F1 Score #

使用sklearn的metrics模块直接计算，支持多分类/二分类场景：

# 生成混淆矩阵
conf_mat = confusion_matrix(true_labels, pred_labels)
print("混淆矩阵：")
print(conf_mat)

# 计算指标，根据任务选择average参数：
# - binary：二分类场景
# - macro：多分类下各类别指标取算术平均
# - weighted：多分类下按样本量加权平均
precision = precision_score(true_labels, pred_labels, average='macro')
recall = recall_score(true_labels, pred_labels, average='macro')
f1 = f1_score(true_labels, pred_labels, average='macro')

print(f"精准率(Precision): {precision:.4f}")
print(f"召回率(Recall): {recall:.4f}")
print(f"F1分数(F1 Score): {f1:.4f}")

3. 绘制ROC曲线与计算AUC #

分二分类和多分类场景分别处理：

二分类场景 #

# 提取正类的预测概率（假设正类对应索引1）
y_score = pred_probs[:, 1]

# 计算ROC曲线参数与AUC值
fpr, tpr, thresholds = roc_curve(true_labels, y_score)
roc_auc = auc(fpr, tpr)

# 绘制ROC曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.4f})')
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假阳性率(FPR)')
plt.ylabel('真阳性率(TPR)')
plt.title('受试者工作特征曲线(ROC)')
plt.legend(loc="lower right")
plt.show()

多分类场景 #

采用One-vs-Rest策略绘制每个类别的ROC曲线：

n_classes = pred_probs.shape[1]
# 将真实标签二值化
true_labels_bin = label_binarize(true_labels, classes=np.arange(n_classes))

# 逐个类别计算ROC与AUC
fpr = dict()
tpr = dict()
roc_auc = dict()
for i in range(n_classes):
    fpr[i], tpr[i], _ = roc_curve(true_labels_bin[:, i], pred_probs[:, i])
    roc_auc[i] = auc(fpr[i], tpr[i])

# 绘制多分类ROC曲线
plt.figure(figsize=(8, 6))
colors = ['aqua', 'darkorange', 'cornflowerblue', 'red', 'green']
for i, color in zip(range(n_classes), colors):
    plt.plot(fpr[i], tpr[i], color=color, lw=2,
             label=f'类别{i}的ROC曲线 (AUC = {roc_auc[i]:.4f})')

plt.plot([0, 1], [0, 1], 'k--', lw=2)
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假阳性率(FPR)')
plt.ylabel('真阳性率(TPR)')
plt.title('多分类受试者工作特征曲线')
plt.legend(loc="lower right")
plt.show()

注意事项 #

• 确保模型处于eval()模式，避免BatchNorm、Dropout等层影响预测结果。
• 所有张量需先转到CPU再转为numpy数组，sklearn不支持CUDA张量计算。
• 二分类场景下，若模型输出未经过sigmoid，需先添加sigmoid激活获取概率；多分类则用softmax。

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