PyTorch中如何将按样本过滤的ImageFolder数据集用于推理？

步骤1：转换待保留类别为对应索引 #

你已经得到了样本量≤20的类别名称列表sampled_classes，首先将其转换为ImageFolder中对应的类别索引：

target_classes_idx = [dataset.class_to_idx[c] for c in sampled_classes]

步骤2：生成过滤后的数据集 #

直接用PyTorch官方的Subset工具筛选符合要求的样本，不需要修改原始ImageFolder结构：

from torch.utils.data import Subset

# 遍历所有样本元信息，收集符合条件的样本索引，无需加载图像效率更高
keep_sample_idx = []
for idx, (_, label) in enumerate(dataset.samples):
    if label in target_classes_idx:
        keep_sample_idx.append(idx)

# 得到过滤后的数据集，完全兼容ImageFolder的输出格式
filtered_dataset = Subset(dataset, keep_sample_idx)

步骤3：构造DataLoader并执行推理 #

用过滤后的数据集构造DataLoader，即可直接代入你提供的推理代码使用：

from torch.utils.data import DataLoader

# batch_size根据你的显存情况调整，推理阶段无需shuffle
data_loader = DataLoader(filtered_dataset, batch_size=32, shuffle=False)

之后直接运行你给出的推理代码即可得到所有小样本类别的预测结果和真实标签。

可选：最终准确率计算 #

推理完成后拼接结果计算整体准确率：

import numpy as np

y_pred = np.concatenate(y_pred_list)
y_true = np.concatenate(y_true_list)
small_class_acc = (y_pred == y_true).mean()
print(f"样本量≤20的类别平均准确率为：{small_class_acc:.4f}")

注意事项 #

• 原始ImageFolder加载时请确认已传入推理所需的transform（如Resize、ToTensor、Normalize等，需和训练阶段的测试预处理逻辑完全一致）
• 如果仅需要计算准确率，不需要调整原始标签映射，直接使用上述Subset方案即可，预测结果和真实标签的定义和原始数据集完全匹配

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