我正在准备数据以训练一个图像识别模型。目前我有一个文件夹（数据集），其中包含多个以标签命名的子文件夹，这些子文件夹内存储着图像。

我想将这个数据集进行分割，以便我有两个主要文件夹，它们拥有相同的子文件夹结构，但这些文件夹内的图像数量按照我偏好的训练/测试分割比例分配，例如训练数据集中的图像占90%，测试数据集中的图像占10%。

我正在寻找最好的方法来分割我的数据。我读到一个建议，说pytorch的torch.utils.Dataset类可能是一种方法，但我似乎无法使其工作以保持文件夹层次结构。

回答：

如果你有一个这样的文件夹结构：

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

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

正如pytorch网站所述

这是一个通用的数据加载器，其中图像按以下方式排列：

root/dog/xxx.pngroot/dog/xxy.pngroot/dog/xxz.pngroot/cat/123.pngroot/cat/nsdf3.pngroot/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))  # 取10%用于测试train_set = torch.utils.data.Subset(dataset, range(test_size, len(dataset)) # 剩余部分用于训练

如果你想对分割进行shuffle，请记住subset类使用索引进行分割。所以你可以先shuffle，再进行分割。做类似这样的事情

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