基于Tensor Dataset预测时，如何保证test_predictions与test_images顺序一致？

问题根源在于你两次独立迭代test_dataset：一次用于模型预测，一次用于提取输入图像。如果测试集开启了shuffle=True（或是默认包含随机操作），两次迭代的样本顺序会被打乱，导致预测结果和输入图像无法对应。

方案1：一次性提取所有样本，保证顺序一致 #

先把测试集的输入图像和目标图像全部提取到内存中，再用提取后的输入做预测，从根源上保证顺序完全匹配：

# 遍历测试集，同步保存输入图像与目标图像
test_images = []
test_targets = []
for x, y in test_dataset:
    test_images.append(x.numpy())
    test_targets.append(y.numpy())

# 用提取好的输入图像生成预测结果
import numpy as np
test_images_np = np.array(test_images)
test_predictions = model.predict(test_images_np)

# 展示第1组对应结果
fig, ax = plt.subplots(1, 2, figsize=(10, 10))
ax[0].imshow(test_predictions[0], cmap='gray')
ax[0].set_title('预测掩码')
ax[1].imshow(test_images[0], cmap='gray')
ax[1].set_title('输入图像')
plt.show()

方案2：批量预测时同步保存输入与结果 #

如果测试集过大，无法一次性加载到内存，可以在批量迭代时同步获取输入和预测结果：

test_dataset_batched = test_dataset.batch(batch_size=1)
test_predictions = []
test_images = []

for x, y in test_dataset_batched:
    # 对当前批量样本做预测
    pred = model.predict(x, verbose=0)
    test_predictions.append(pred[0])
    test_images.append(x[0].numpy())

# 展示对应结果
fig, ax = plt.subplots(1, 2, figsize=(10, 10))
ax[0].imshow(test_predictions[0], cmap='gray')
ax[1].imshow(test_images[0], cmap='gray')
plt.show()

关键注意事项 #

• 测试集必须设置shuffle=False，避免迭代时样本顺序随机变化。
• 永远不要分开两次迭代测试集，始终从同一个迭代流程中同步获取输入和预测结果，才能保证一一对应。

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