TFLite评估报错：无法将(128,128,3)形状数组广播到形状(1,)如何解决

报错根因 #

这个报错由两处逻辑不匹配共同导致：

1. 你参考的量化示例用的是MNIST手写数字数据集，直接加载的test_images是形状为(样本数, 28,28)的numpy数组，遍历即可直接拿到单张二维图像；但你使用的dataframe iterator（如Keras的flow_from_dataframe生成的迭代器）每次返回的是**(批次图像张量, 批次标签张量)的元组**，不是单张图像，直接遍历拿到的元素本身已经带batch维度，照搬示例的expand_dims逻辑会追加错误维度，触发广播错误。
2. 你的输入是128128尺寸的3通道RGB图像，和示例里2828的单通道MNIST图像维度完全不同，且没有针对自己的TFLite模型输入要求做适配——如果是全整数量化的模型，输入要求int8类型，直接传float32数组也会触发类型/维度匹配错误。

修复步骤 #

• 初始化测试集dataframe iterator时，将shuffle参数设为False，保证预测结果顺序和test_labels顺序一一对应，避免准确率计算错误；调试阶段可以先把batch_size设为1，简化单张图像处理逻辑。
• 推理前先读取TFLite模型的输入详情，确认输入形状、数据类型、量化参数，不要硬编码维度。
• 修正迭代器遍历逻辑，从返回的元组中拆分出图像数据，按模型要求做预处理、维度补全、量化适配后再传入推理。
• 如果是int8量化模型，需要对输入做量化转换、对输出做反量化，否则推理结果会完全失效。

修复后可运行代码 #

def evaluate_model(interpreter):
    input_details = interpreter.get_input_details()[0]
    output_details = interpreter.get_output_details()[0]
    input_index = input_details["index"]
    output_index = output_details["index"]
    target_input_shape = input_details["shape"]
    target_input_dtype = input_details["dtype"]

    prediction_results = []
    # 重置迭代器，避免之前的遍历导致偏移
    test_images.reset()
    total_test_samples = test_images.samples

    for _ in range(total_test_samples):
        # 迭代器每次返回(批次图像, 批次标签)，batch_size=1时单次取1张图
        image_batch, _ = next(test_images)
        # 取出批次内的单张图像，形状为(128,128,3)
        single_image = image_batch[0]

        # 适配int8量化模型的输入要求
        if target_input_dtype == np.int8:
            input_scale, input_zero_point = input_details["quantization"]
            single_image = single_image / input_scale + input_zero_point
            single_image = single_image.astype(target_input_dtype)
        
        # 补充batch维度，匹配模型输入形状 (1,128,128,3)
        input_tensor = np.expand_dims(single_image, axis=0)
        # 提前校验形状，避免广播错误
        assert input_tensor.shape == tuple(target_input_shape), \
            f"输入形状不匹配，预期{target_input_shape}，实际{input_tensor.shape}"
        
        interpreter.set_tensor(input_index, input_tensor)
        interpreter.invoke()

        # 读取输出并处理int8反量化
        output_tensor = interpreter.get_tensor(output_index)
        if output_details["dtype"] == np.int8:
            output_scale, output_zero_point = output_details["quantization"]
            output_tensor = (output_tensor - output_zero_point) * output_scale
        
        pred_label = np.argmax(output_tensor[0])
        prediction_results.append(pred_label)

        if len(prediction_results) % 1000 == 0:
            print(f'Evaluated on {len(prediction_results)} results so far.')

    # 计算准确率
    prediction_results = np.array(prediction_results)
    accuracy = (prediction_results == test_labels).mean()
    return accuracy

补充说明 #

如果你不想修改迭代器的batch_size，只需要在拿到image_batch后加一层循环，遍历批次内的每一张图逐一推理即可，逻辑和单张处理完全一致。

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