使用model.fit_generator预测图像出现维度错误，如何将预测结果保存到文件夹

问题根因 #

Keras卷积层要求的输入为4维张量，格式为(batch_size, height, width, channels)，你传入model.predict()的是单张图像，形状为(224,224,3)，缺少了第一个批量维度，所以触发了维度校验失败。

问题出在循环预测的这一行代码：

predicted[i][j]= model.predict(data_list[i] [j])

data_list[i]是单个批次的图像数据，形状为(batch_size, 224,224,3)，取data_list[i][j]相当于单独拿出了批次里的第j张图，直接丢失了批量维度。

修复方案 #

方案1：最小改动，给单张图像增加批量维度 #

导入numpy后，在预测前给输入增加第0维即可：

import numpy as np

# 原有生成data_list的代码保持不变
for i in range(batch_index):
    # 适配最后一个批次可能不足batch_size的情况，避免越界
    for j in range(len(data_list[i])):
        # 给单张图增加批量维度，符合模型输入要求
        input_img = np.expand_dims(data_list[i][j], axis=0)
        # 取预测结果的第一个元素，去掉多余的批量维度再保存
        predicted_single = model.predict(input_img, verbose=0)[0]
        plt.imsave(os.path.join(path,f'new_image{i}{j}_AE_.tiff'), predicted_single)

方案2：批次级预测，运行效率更高 #

不需要逐张循环预测，直接对整个批次做预测，减少计算和IO开销：

for i in range(batch_index):
    # 直接预测整个批次，一次得到所有样本的结果
    predicted_batch = model.predict(data_list[i], verbose=0)
    for j in range(len(predicted_batch)):
        plt.imsave(os.path.join(path,f'new_image{i}{j}_AE_.tiff'), predicted_batch[j])

内容的提问来源于stack exchange，提问作者T. Elhassan