按文件数排序目录构建斯坦福犬不均衡数据集的PyTorch使用问题

原生ImageFolder默认会自动扫描根目录下所有子目录加载数据，不需要直接传入自定义目录列表，你可以通过先全量加载再按规则筛选样本的方式实现需求，不需要修改原目录结构或重命名文件，方案稳定可控：

步骤1：处理已排序的目录数据 #

假设你已经拿到按文件数降序排列的目录字典class_count_dict（key为类别目录路径，value为对应目录下的文件数），先将其转为有序列表：

# 按文件数降序排序类别
sorted_class_dirs = sorted(class_count_dict.items(), key=lambda x: x[1], reverse=True)

步骤2：全量加载数据集后按规则筛选样本 #

通过Subset工具按需筛选样本，不需要修改原始数据集文件：

import os
import random
from torchvision.datasets import ImageFolder
from torch.utils.data import Subset

# 先加载全量数据集
full_dataset = ImageFolder(root="替换为你的数据集根目录路径")

# 按类别存储所有样本的索引
class_idx_map = {}
for sample_idx, (img_path, label) in enumerate(full_dataset.samples):
    if label not in class_idx_map:
        class_idx_map[label] = []
    class_idx_map[label].append(sample_idx)

# 按规则筛选需要的样本索引
selected_indices = []
for rank, (class_dir, _) in enumerate(sorted_class_dirs):
    # 计算当前类别应保留的样本数
    keep_count = 252 - 2 * rank
    keep_count = max(keep_count, 1) # 避免保留数为负
    # 匹配当前类别对应的标签
    class_name = os.path.basename(class_dir)
    label = full_dataset.class_to_idx[class_name]
    # 随机采样指定数量的样本，需要固定结果可提前设置random.seed
    selected = random.sample(class_idx_map[label], keep_count)
    selected_indices.extend(selected)

# 生成最终的不均衡数据集
imbalanced_dog_dataset = Subset(full_dataset, selected_indices)

步骤3：验证结果 #

可以直接统计输出每个类别的样本数确认符合要求：

from collections import Counter
label_count = Counter([imbalanced_dog_dataset[i][1] for i in range(len(imbalanced_dog_dataset))])
print(sorted(label_count.values(), reverse=True))
# 正常输出应为 [252, 250, 248, 246...] 符合你的预期

如果不需要随机采样，要固定取每个类别前N张样本，把random.sample替换为切片class_idx_map[label][:keep_count]即可。

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