HOWTO · Python
在 Python 中绘制混淆矩阵
学习在 Python 中用 Matplotlib、Seaborn 或 Pretty Confusion Matrix 绘制并标注混淆矩阵。
本页内容
要在 Python 中绘制混淆矩阵，先统计每个真实类别与预测类别组合的次数，然后用 Matplotlib 或 Seaborn 显示这些计数。本文还介绍了基于包的 Pretty Confusion Matrix 选项。行表示真实类别，列表示预测类别。
在 Python 中使用 Matplotlib 绘制混淆矩阵
混淆矩阵存储每个真实类别与预测类别组合的计数。只定义一次类别顺序，使矩阵的行和列与坐标轴标签相符。在这个例子中，y_true 包含每个样本的观测类别，y_pred 包含相应的模型预测。第一行和第一列代表 bat；第二行和第二列代表 ball。下面这个可运行的示例只使用 Python 标准库来构建计数矩阵。
运行此示例：
classes = ["bat", "ball"]
y_true = ["bat", "ball", "ball", "bat", "bat", "bat"]
y_pred = ["bat", "bat", "ball", "ball", "bat", "bat"]
matrix = [
[sum(actual == row and predicted == column
for actual, predicted in zip(y_true, y_pred))
for column in classes]
for row in classes
]
print(matrix)
[[3, 1], [1, 1]]
下面这个完整的 Matplotlib 程序会把完成的图保存到 confusion-matrix-matplotlib.png，因此无论有没有显示器都能同样运行；在交互式会话中，你可以把 savefig 一行换成 plt.show()。
import matplotlib.pyplot as plt
classes = ["bat", "ball"]
y_true = ["bat", "ball", "ball", "bat", "bat", "bat"]
y_pred = ["bat", "bat", "ball", "ball", "bat", "bat"]
matrix = [
[sum(actual == row and predicted == column
for actual, predicted in zip(y_true, y_pred))
for column in classes]
for row in classes
]
fig, ax = plt.subplots(figsize=(5, 4))
image = ax.imshow(matrix, cmap="Blues", origin="upper")
fig.colorbar(image, ax=ax, label="Count")
for row in range(len(classes)):
for column in range(len(classes)):
value = matrix[row][column]
ax.text(column, row, str(value), ha="center", va="center",
color="white" if value > 1.5 else "black")
ax.set_xticks(range(len(classes)))
ax.set_yticks(range(len(classes)))
ax.set_xticklabels(classes)
ax.set_yticklabels(classes)
ax.set_xlabel("Predicted class")
ax.set_ylabel("Actual class")
ax.set_title("Confusion Matrix")
fig.tight_layout()
fig.savefig("confusion-matrix-matplotlib.png", bbox_inches="tight")
矩阵计算得到 [[3, 1], [1, 1]]：三只真实的 bat 被预测为 bat，一只真实的 bat 被预测为 ball，一只真实的 ball 被预测为 bat，一只真实的 ball 被预测为 ball。对角线单元格统计正确预测；非对角线单元格统计错误分类。这些是计数，不是百分比。
图中的行是真实标签，列是预测标签。标注循环读取 matrix[row][column]，但传给 ax.text 的是 (column, row)，因为 x 是水平方向，y 是垂直方向。对于像 [[3, 2], [1, 4]] 这样的不对称矩阵，这会把 2 放在 x=1、y=0，把 1 放在 x=0、y=1；交换这些坐标会标错单元格。origin="upper" 让第一个真实类别保持在顶部。
改编这个程序时，提供等长的标签列表，把相应的观测保持在相同位置，并把每个观测到的类别都包含在 classes 中。否则，zip 会截断不匹配的项，或者推导式会省略你没有声明的类别。
在 Python 中使用 Seaborn 绘制混淆矩阵
Seaborn 的 heatmap 会替你处理单元格标注。这个独立程序定义了自己的数据和矩阵，因此不依赖于先运行前一节。它需要 Seaborn 和 Matplotlib，并把图保存到 confusion-matrix-seaborn.png。
import matplotlib.pyplot as plt
import seaborn as sns
classes = ["bat", "ball"]
y_true = ["bat", "ball", "ball", "bat", "bat", "bat"]
y_pred = ["bat", "bat", "ball", "ball", "bat", "bat"]
matrix = [
[sum(actual == row and predicted == column
for actual, predicted in zip(y_true, y_pred))
for column in classes]
for row in classes
]
fig, ax = plt.subplots(figsize=(5, 4))
sns.heatmap(matrix, annot=True, fmt="d", cmap="Blues", square=True,
xticklabels=classes, yticklabels=classes,
cbar_kws={"label": "Count"}, ax=ax)
ax.set_xlabel("Predicted class")
ax.set_ylabel("Actual class")
ax.set_title("Confusion Matrix")
ax.tick_params(axis="y", labelrotation=0)
fig.tight_layout()
fig.savefig("confusion-matrix-seaborn.png", bbox_inches="tight")
在这里，annot=True 把数值写入单元格，fmt="d" 把它们格式化为整数。计数计算同样得到 [[3, 1], [1, 1]]。两个刻度标签列表都遵循声明的顺序，先 bat 后 ball，而不是无关的布尔标签。真实类别沿垂直轴排列；预测类别沿水平轴排列。使用你自己的数据时，要把配对的观测和类别列表一起替换。
在 Python 中使用 Pretty Confusion Matrix 绘制混淆矩阵
Pretty Confusion Matrix 提供 pp_matrix_from_data 入口，用配对的真实标签和预测标签绘图。该包的 PyPI 项目页面提供其 API 示例、安装指导和依赖元数据。发行版名称是 pretty-confusion-matrix，而 Python 导入使用下划线：pretty_confusion_matrix。
下面这个完整示例会把该包的图保存到 confusion-matrix-pretty.png。它已按原文针对该包及其依赖的当前版本执行过；搭建你自己的环境时，请查阅该项目的依赖要求。
import matplotlib.pyplot as plt
from pretty_confusion_matrix import pp_matrix_from_data
classes = ["bat", "ball"]
y_true = ["bat", "ball", "ball", "bat", "bat", "bat"]
y_pred = ["bat", "bat", "ball", "ball", "bat", "bat"]
# Numeric IDs make the package's sorted class order match the display names.
class_ids = {label: index for index, label in enumerate(classes)}
true_values = [class_ids[label] for label in y_true]
predicted_values = [class_ids[label] for label in y_pred]
pp_matrix_from_data(true_values, predicted_values,
columns=classes, pred_val_axis="col")
plt.savefig("confusion-matrix-pretty.png", bbox_inches="tight")
true_values 包含真实类别的 ID，predicted_values 包含相应的预测。把 bat 编码为 0、ball 编码为 1，使它们排序后的顺序与 columns=classes 一致；columns 参数提供显示名称，而不是对观测重新排序。pred_val_axis="col" 选择预测类别为列，真实类别为行。底层的类别计数是 [[3, 1], [1, 1]]；该包的呈现还可以包含汇总总计，因此它的布局与上面两张普通热力图不同。这里没有声称任何额外的指标结果。
这个程序生成的图是：
