如何在PyTorch中使用random_split()处理含多子文件夹的数据集（无需自定义Dataset）

嘿，我看到你遇到的问题了，咱们先拆解一下代码里的错误，再给你一个完美适配的解决方案：

首先说你代码里的核心问题 #

• 图片数量统计错了：你用len(os.path.join(path, file))算的是文件夹路径的字符长度，不是文件夹里图片的数量！这直接导致划分的长度完全不对，触发了那个“长度总和不匹配”的错误。
• random_split用错了：这个函数的第一个参数必须是PyTorch的Dataset对象，你传了个文件夹名字字符串，这完全不符合API要求呀。
• 还有小拼写错误：比如files.endswith应该是file.endswith，tsfnames.appned少写了个d，而且train_ratio、val_ratio还有trfnames这些变量都没提前初始化。

无需自定义Dataset的正确做法 #

你的数据集结构正好适配PyTorch内置的ImageFolder，它会自动处理子文件夹对应的类别，生成标准的Dataset，然后我们直接对这个整体Dataset做划分就行，不用手动遍历每个子文件夹：

完整可运行代码 #

import os
from torch.utils.data import random_split
from torchvision.datasets import ImageFolder

# 先定义好参数，你可以根据需求调整比例
path = '/Users/francesca/Documents/Tesi/imagenet/'
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

# 用ImageFolder加载整个数据集，自动识别子文件夹为类别
full_dataset = ImageFolder(root=path)

# 计算各集合的样本数量
total_samples = len(full_dataset)
train_len = int(total_samples * train_ratio)
val_len = int(total_samples * val_ratio)
test_len = total_samples - train_len - val_len

# 执行划分，得到三个子集
train_set, val_set, test_set = random_split(full_dataset, lengths=[train_len, val_len, test_len])

# 如果需要获取每个子集里的文件名和对应标签，可以用这个辅助函数
def get_subset_samples(subset):
    samples_info = []
    # 子集的indices存储了它在原数据集里的索引
    for idx in subset.indices:
        img_path, label = subset.dataset.samples[idx]
        # 这里可以根据需要存完整路径或者只存文件名
        samples_info.append((os.path.basename(img_path), subset.dataset.classes[label]))
    return samples_info

# 查看结果
train_samples = get_subset_samples(train_set)
val_samples = get_subset_samples(val_set)
test_samples = get_subset_samples(test_set)

print(f"训练集共 {len(train_samples)} 个样本")
print(f"验证集共 {len(val_samples)} 个样本")
print(f"测试集共 {len(test_samples)} 个样本")

关键细节说明 #

• ImageFolder的优势：它会自动遍历主文件夹下的所有非隐藏子文件夹，把每个子文件夹的名字作为类别标签，子文件夹里的图片自动关联到对应类别，完全不需要你手动处理分类逻辑。
• random_split的正确用法：传入完整的Dataset和各部分的长度列表，它会返回三个Subset对象，这些对象可以直接放进DataLoader里用于训练、验证和测试。
• 获取文件名/标签：Subset对象通过.indices保存了它在原Dataset中的索引，我们可以通过原Dataset的.samples属性拿到对应的图片路径和标签值，再通过.classes把标签值转成类别名称。

可选：按类别分层划分（避免类别分布不均） #

如果你的数据集里不同类别的样本数量差异很大，或者希望每个类别都按比例分到三个集合里（防止某个类别在测试集里完全没有），可以用下面的方法，每个类别单独划分后再合并：

import os
from torch.utils.data import Subset, random_split
from torchvision.datasets import ImageFolder

path = '/Users/francesca/Documents/Tesi/imagenet/'
train_ratio = 0.7
val_ratio = 0.2

full_dataset = ImageFolder(root=path)
# 先把每个类别的样本索引分组
class_indices = {cls: [] for cls in full_dataset.classes}
for idx, (_, label) in enumerate(full_dataset.samples):
    class_name = full_dataset.classes[label]
    class_indices[class_name].append(idx)

train_indices = []
val_indices = []
test_indices = []

# 遍历每个类别，单独划分
for cls_idx_list in class_indices.values():
    cls_total = len(cls_idx_list)
    tr_len = int(cls_total * train_ratio)
    vl_len = int(cls_total * val_ratio)
    ts_len = cls_total - tr_len - vl_len
    # 对当前类别的索引列表做划分
    tr_idx, vl_idx, ts_idx = random_split(cls_idx_list, lengths=[tr_len, vl_len, ts_len])
    train_indices.extend(tr_idx)
    val_indices.extend(vl_idx)
    test_indices.extend(ts_idx)

# 构建最终的子集
train_set = Subset(full_dataset, train_indices)
val_set = Subset(full_dataset, val_indices)
test_set = Subset(full_dataset, test_indices)

这样每个类别都会按比例出现在三个集合里，适合类别分布不均衡的场景。

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