如何用PyTorch将按类别分文件夹的云图像数据集拆分为训练/验证/测试集

直接用ImageFolder加载完整数据集后，结合PyTorch的random_split就能完成拆分，以下是具体实现步骤：

1. 导入依赖库 #

import torch
from torchvision.datasets import ImageFolder
from torch.utils.data import random_split, DataLoader
from torchvision import transforms

2. 定义数据预处理变换 #

根据任务需求定义图像预处理操作，比如统一尺寸、转张量、归一化：

transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 统一图像尺寸
    transforms.ToTensor(),          # 转为PyTorch张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

3. 用ImageFolder加载完整数据集 #

指定数据集根目录和预处理变换，ImageFolder会自动按子文件夹划分类别：

full_dataset = ImageFolder(root='./dataset_folder', transform=transform)

4. 设置拆分比例并计算子集大小 #

比如按70%训练集、20%验证集、10%测试集拆分，根据总数据量计算各子集的样本数：

total_size = len(full_dataset)
train_size = int(0.7 * total_size)
val_size = int(0.2 * total_size)
test_size = total_size - train_size - val_size

5. 拆分数据集 #

使用random_split拆分，设置随机种子保证拆分结果可复现：

# 设置随机种子
torch.manual_seed(42)

train_dataset, val_dataset, test_dataset = random_split(full_dataset, [train_size, val_size, test_size])

6. 创建数据加载器 #

为三个子集分别创建DataLoader，方便后续模型训练时批量加载数据：

batch_size = 32

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

验证拆分结果 #

可以打印各子集的样本数量和类别信息，确认拆分是否正确：

print(f"训练集样本数: {len(train_dataset)}")
print(f"验证集样本数: {len(val_dataset)}")
print(f"测试集样本数: {len(test_dataset)}")
print(f"类别列表: {full_dataset.classes}")

注意事项 #

• 如果需要按类别比例拆分（保证每个子集的类别分布和原数据集一致），可以先通过sklearn.model_selection.train_test_split拆分索引，再基于索引构建子集；也可以自定义采样器实现。
• 随机种子要在拆分前设置，避免每次运行拆分结果不同。

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