在Pytorch中，文件夹名称可用作标签的多类图像数据集的馈送方法？

Question

我想在Pytorch中提供多类图像数据集，在数据集的主文件夹中，我还有15个不同名称的文件夹，我想使用文件夹名称作为标签。 例如，一个文件夹名为Aeroplanes，包含图像（1245个图像），另一个文件夹名为Cars，包含汽车图像（997），同样，每个文件夹都有不同数量的图像。 现在我想加载它们来训练我的模型并对其进行测试，但我没有用于训练和测试的单独文件夹。我希望使用文件夹名称作为标签，还希望将数据集按相等的比例拆分为培训和测试 在这种情况下，我们将感谢您的指导。谢谢

Answer 1

要将数据集拆分为训练和测试数据集，可以使用random_split函数：

import torch
from torchvision import datasets, transforms
from torch.utils import data
import numpy as np

dataset = datasets.ImageFolder('path_to_dataset', transform=transforms.ToTensor())

lengths = [int(np.ceil(0.5*len(dataset))),
           int(np.floor(0.5*len(dataset)))]
train_set, test_set = data.random_split(dataset, lengths)

train_dataloader = data.DataLoader(train_set, batch_size=...)
test_dataloader = data.DataLoader(test_set, batch_size=...)

如果您想在列车和测试数据集上执行单独的转换，请查看此处：How to use different data augmentation for Subsets in PyTorch