图像分类代码添加混淆矩阵输出功能

在图像分类代码中添加混淆矩阵输出功能

为了更好地评估图像分类模型的性能，我们可以添加混淆矩阵输出功能。以下步骤将在给定的代码基础上添加混淆矩阵输出功能，并提供具体行数的修改或新增建议。

导入必要的库

import numpy as np
from sklearn.metrics import confusion_matrix

在测试循环之前初始化一个空的混淆矩阵

# 新增以下行
num_classes = len(class_names)
confusion_mat = np.zeros((num_classes, num_classes))

在测试循环中，每次预测一个图像后，更新混淆矩阵

# 修改以下行
outputs = model(images)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu().numpy()
labels = labels.cpu().numpy()

# 新增以下行
confusion_mat += confusion_matrix(labels, predicted, labels=range(num_classes))

在测试循环结束后，打印混淆矩阵

# 新增以下行
print("\nConfusion Matrix:")
print(confusion_mat)

注意：

以上修改假设你已经导入了必要的库，并且代码中的其他部分保持不变。
混淆矩阵的打印格式可以根据需要进行调整。

通过以上步骤，你就可以在图像分类代码中成功添加混淆矩阵输出功能，并更直观地评估模型的分类效果。