计算预训练CNN模型混淆矩阵时样本数量不一致报错求解

这个报错的核心是真实标签和预测结果的样本数量不匹配：你的测试集真实标签只有16979个，但生成的预测结果有271664个，导致混淆矩阵无法计算，具体成因有三个：

• 预训练模型加载不规范：加载MobileNet时没有设置include_top=False，保留了原模型针对ImageNet数据集设计的1000类分类头，后续叠加自定义分类层后，模型输出维度和逻辑异常，导致预测结果数量不符合预期。
• 预测步数计算错误：手动用整除//计算测试步数，当测试集总数无法被batch_size整除时，要么会漏掉最后一个不完整批次的样本，要么旧版本Keras的生成器会循环输出数据，导致预测样本数远大于真实样本数。
• 使用弃用接口：predict_generator已经在新版TensorFlow中被弃用，接口逻辑与当前主流的predict接口存在差异，容易触发意料之外的错误。

第一步：修正预训练模型加载 #

加载预训练模型时去掉原模型的顶层分类头，指定匹配的输入形状，替换代码如下：

# 原代码：basemodel = tf.keras.applications.mobilenet.MobileNet()
basemodel = tf.keras.applications.mobilenet.MobileNet(include_top=False, input_shape=(224, 224, 3))

第二步：修正预测逻辑 #

删除手动计算测试步数的逻辑，直接使用model.predict接口自动遍历所有测试样本，替换代码如下：

# 删掉原有的STEP_SIZE_TEST计算逻辑
test_data.reset()
y_pred = model.predict(test_data, verbose=1)
predicted_class = np.argmax(y_pred, axis=1)

model.predict传入数据生成器时会自动匹配所有样本，不会出现多生成或者漏生成样本的问题。

第三步（兼容旧版本TensorFlow可选） #

如果你使用的TensorFlow版本必须指定预测步数，将步数计算改为向上取整，同时截取和真实标签数量一致的预测结果：

import numpy as np
STEP_SIZE_TEST = int(np.ceil(test_data.n / test_data.batch_size))
test_data.reset()
y_pred = model.predict(test_data, steps=STEP_SIZE_TEST, verbose=1)
predicted_class = np.argmax(y_pred, axis=1)[:test_data.n]

完成以上修改后，预测结果的样本数量会和真实标签完全匹配，混淆矩阵即可正常计算。

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