为嵌套在多个文件夹中的数据创建一个训练、测试拆分



我正在为训练图像识别模型准备数据。我目前有一个文件夹(数据集(,其中包含多个带有标签名称的文件夹,这些文件夹中有图像。

我想以某种方式拆分这个数据集,这样我就有两个具有相同子文件夹的主文件夹,但这些文件夹中的图像数量要根据首选的训练/测试拆分,例如训练数据集中90%的图像和测试数据集中10%的图像。

我正在努力寻找如何拆分数据的最佳方法。我读到一个建议,pytorch torch.utils.Dataset类可能是一种方法,但我似乎无法让它发挥作用,因为它可以保留文件夹层次结构。

如果您有这样的文件夹结构:

folder
│     
│
└───class1
│   │   file011
│   │   file012
│   
└───class2
│   file021
│   file022

您可以简单地使用类torchvision.datasets.ImageFolder

如pytorch 网站所述

一个通用的数据加载器,其中图像以这种方式排列:

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

然后,在创建ImageFolder实例后,例如

dataset = torchvision.datasets.Imagefolder(YOUR_PATH, ...)

你可以这样拆分:

test_size = 0.1 * len(dataset)
test_set = torch.utils.data.Subset(dataset, range(test_size))  # take 10% for test
train_set = torch.utils.data.Subset(dataset, range(test_size, len(dataset)) # the last part for train

如果要对拆分进行混洗,请记住类子集使用拆分的索引。所以你可以洗牌,然后把它们分开。做类似的事情

indexes = shuffle(range(len(dataset)))
indexes_train = indexes[:int(len(dataset)*0.9)]
indexes_test = = indexes[int(len(dataset)*0.9):]

最新更新