如何使用 torchvision.datasets.Imagefolder 将数据拆分为训练集和测试集？

Question

在我的自定义数据集中，一种图像在一个文件夹中，torchvision.datasets.Imagefolder 可以处理，但如何将数据集拆分为训练和测试？

Answer 1

您可以根据示例的索引torch.utils.data.Subset将ImageFolder数据集拆分为训练和测试。
例如：

orig_set = torchvision.datasets.Imagefolder(...)  # your dataset
n = len(orig_set)  # total number of examples
n_test = int(0.1 * n)  # take ~10% for test
test_set = torch.utils.data.Subset(orig_set, range(n_test))  # take first 10%
train_set = torch.utils.data.Subset(orig_set, range(n_test, n))  # take the rest