使用datasets.ImageFolder()划分数据集并打乱顺序（简单易懂）

本文链接：https://blog.csdn.net/TycoonL/article/details/125667592

这篇博客介绍了如何使用Python的torch库来划分和打乱数据集，以避免标签扎堆现象。首先，定义了数据路径、转换操作、测试集比例和批处理大小。接着，通过random.sample选择测试集样本，然后利用Subset创建训练集和测试集。使用DataLoader加载数据时，训练集启用shuffle，而测试集则不启用。博主指出，直接在DataLoader中打乱已经划分的数据集可能无效，因为数据集标签的连续性可能导致划分后的数据集中标签分布不均。通过这段代码，可以确保训练集和测试集的合理分布。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、代码

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import random

path = 
transforms=
proportion=0.1 #测试集比例
batch_size=32

data = datasets.ImageFolder(path,transforms)
n = len(data)  #数据集总数
n_test = random.sample(range(1, n), int(proportion * n))  #按比例取随机数列表

test_set = torch.utils.data.Subset(data, n_test)  #按照随机数列表取测试集
train_set = torch.utils.data.Subset(data,list(set(range(1, n)).difference(set(n_test))))  #测试集剩下作为训练集

data_train = DataLoader(train_set, batch_size=batch_size, shuffle=True)
data_test=DataLoader(test_set, batch_size=batch_size, shuffle=False)

#输出筛选的训练集labels
for batch_idex, (data, targets) in enumerate(data_test):
    print(batch_idex,targets)