代码有些参考了其他博客
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import itertools
def plot_confusion_matrix(cm, classes, normalize=False, cmap=plt.cm.Blues):
"""
This function prints and plots the confusion matrix.
Normalization can be applied by setting `normalize=True`.
Input
- cm : 计算出的混淆矩阵的值
- classes : 混淆矩阵中每一行每一列对应的列
- normalize : True:显示百分比, False:显示个数
"""
if normalize:
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
title = 'Normalized confusion matrix'
cb_title = 'Number of recordings (normalized)'
else:
title = 'Confusion matrix'
cb_title = 'Number of recordings'
print(title)
print(cm)
plt.imshow(cm, interpolation='nearest', cmap=cmap)
plt.title(title)
cb = plt.colorbar()
cb.set_label(cb_title)
tick_marks = np.arange(len(classes))
plt.xticks(tick_marks, classes, rotation=90)
plt.yticks(tick_marks, classes)
plt.axis("equal")
ax = plt.gca() # 获得当前axis
left, right = plt.xlim() # 获得x轴最大最小值
ax.spines['left'].set_position(('data', left))
ax.spines['right'].set_position(('data', right))
for edge_i in ['top', 'bottom', 'right', 'left']:
ax.spines[edge_i].set_edgecolor("white")
thresh = cm.max() / 2.
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
num = np.round(cm[i, j], 2) if normalize else int(cm[i, j]) #'{:.2f}'.format(cm[i, j])
plt.text(j, i, num,
verticalalignment='center',
horizontalalignment="center",
color="white" if num > thresh else "black")
plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label')
plt.show()
测试
from sklearn.metrics import confusion_matrix
# pytorch 模型输出的gpu tensor转换到cpu上
y_true = true_all.cpu().numpy()
y_pred = pred_all.cpu().numpy()
cm = confusion_matrix(y_true, y_pred)
# 输出数量的混淆矩阵
plot_confusion_matrix(cm, classes, normalize=False)
# 输出数量归一化后的混淆矩阵
plot_confusion_matrix(cm, classes, normalize=True)